Commit 4283949a authored by Bixia Zheng's avatar Bixia Zheng Committed by TensorFlower Gardener
Browse files

Allow the use of 32 bit integer type for loop index and tensor element index.

The GPU LLVM IR generator currently uses 64 bit integer type for arithmetic
operations related to loop index and tensor element index and relies on LLVM
optimization to narrow the operations to 32 bit integer type. There are
situations whether LLVM optimization fail to perform such an optimization, see
LLVM D46760 for more detail.

This change modifies the XLA LLVM IR code generation infrastructure to support
the use of 32 bit integer type for loop index and tensor element index as
follows:
 .Extends the loop emitter interface in ParallelLoopEmitter and ForLoopNest to
  allow users to specify the loop index type.
 .Modifies the tensor access interface in IrArray::Index interface to record
  the llvm type for the index when an object is constructed. This index type
  is usually propagated from a loop index type.
 .Modifies kernel_support_library to retrieve the loop index type from the
  input llvm::Value.
 .Modifies elemental_ir_emitter to retrieve the data type from the input
  IrArray::Index and use it tensor offset expression.

This change also modifies the emission of the fusion kernel, the row and
scalar reduction kernel and SelectAndScatter kernel to use 32 bit integer type
for index calculation when the size of the launch dimension and the size of
tensors used in the kernel are within the range of 32 bit integer
representation.

PiperOrigin-RevId: 201303468
parent 7c754a6d
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment