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
Loading
Please sign in to comment