[mlir][LLVM] NFC - Remove createIndexConstant method

This revision removes the createIndexConstant method, which implicitly creates constants of the
getIndexType type and updates all uses to the more explicit createIndexAttrConstant which requires
an explicit Type parameter.

This is an NFC step towards entangling index type conversion in LLVM lowering.

The selection of which index type to use requires finer granularity than the existing
implementations which all rely on pass level flags and end up in mismatches, especially on GPUs
with multiple address spaces of different capacities.

This revision also includes an NFC fix to MemRefToLLVM.cpp that prevents a crash in cases where
an integer memory space cannot be derived for a MemRef.

Differential Revision: https://reviews.llvm.org/D156854

Author: nicolasvasilache

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

@@ -65,10 +65,6 @@ class ConvertToLLVMPattern : public ConversionPattern {
   static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
                                        Type resultType, int64_t value);
 
-  /// Create an LLVM dialect operation defining the given index constant.
-  Value createIndexConstant(ConversionPatternRewriter &builder, Location loc,
-                            uint64_t value) const;
-
   // This is a strided getElementPtr variant that linearizes subscripts as:
   //   `base_offset + index_0 * stride_0 + ... + index_n * stride_n`.
   Value getStridedElementPtr(Location loc, MemRefType type, Value memRefDesc,
@@ -155,9 +151,9 @@ class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
                ConversionPatternRewriter &rewriter) const final {
     if constexpr (SourceOp::hasProperties())
       return rewrite(cast<SourceOp>(op),
-              OpAdaptor(operands, op->getDiscardableAttrDictionary(),
-                        cast<SourceOp>(op).getProperties()),
-              rewriter);
+                     OpAdaptor(operands, op->getDiscardableAttrDictionary(),
+                               cast<SourceOp>(op).getProperties()),
+                     rewriter);
     rewrite(cast<SourceOp>(op),
             OpAdaptor(operands, op->getDiscardableAttrDictionary()), rewriter);
   }

mlir/include/mlir/Conversion/MemRefToLLVM/AllocLikeConversion.h

@@ -15,7 +15,7 @@ namespace mlir {
 
 /// Lowering for memory allocation ops.
 struct AllocationOpLLVMLowering : public ConvertToLLVMPattern {
-  using ConvertToLLVMPattern::createIndexConstant;
+  using ConvertToLLVMPattern::createIndexAttrConstant;
   using ConvertToLLVMPattern::getIndexType;
   using ConvertToLLVMPattern::getVoidPtrType;
 
@@ -43,7 +43,9 @@ struct AllocationOpLLVMLowering : public ConvertToLLVMPattern {
     MemRefType memRefType = op.getType();
     Value alignment;
     if (auto alignmentAttr = op.getAlignment()) {
-      alignment = createIndexConstant(rewriter, loc, *alignmentAttr);
+      Type indexType = getIndexType();
+      alignment =
+          createIndexAttrConstant(rewriter, loc, indexType, *alignmentAttr);
     } else if (!memRefType.getElementType().isSignlessIntOrIndexOrFloat()) {
       // In the case where no alignment is specified, we may want to override
       // `malloc's` behavior. `malloc` typically aligns at the size of the

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -168,7 +168,7 @@ struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
     Value lowHalf = rewriter.create<LLVM::TruncOp>(loc, llvmI32, ptrAsInt);
     resource = rewriter.create<LLVM::InsertElementOp>(
         loc, llvm4xI32, resource, lowHalf,
-        this->createIndexConstant(rewriter, loc, 0));
+        this->createIndexAttrConstant(rewriter, loc, this->getIndexType(), 0));
 
     // Bits 48-63 are used both for the stride of the buffer and (on gfx10) for
     // enabling swizzling. Prevent the high bits of pointers from accidentally
@@ -180,7 +180,7 @@ struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
         createI32Constant(rewriter, loc, 0x0000ffff));
     resource = rewriter.create<LLVM::InsertElementOp>(
         loc, llvm4xI32, resource, highHalfTruncated,
-        this->createIndexConstant(rewriter, loc, 1));
+        this->createIndexAttrConstant(rewriter, loc, this->getIndexType(), 1));
 
     Value numRecords;
     if (memrefType.hasStaticShape()) {
@@ -202,7 +202,7 @@ struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
     }
     resource = rewriter.create<LLVM::InsertElementOp>(
         loc, llvm4xI32, resource, numRecords,
-        this->createIndexConstant(rewriter, loc, 2));
+        this->createIndexAttrConstant(rewriter, loc, this->getIndexType(), 2));
 
     // Final word:
     // bits 0-11: dst sel, ignored by these intrinsics
@@ -227,7 +227,7 @@ struct RawBufferOpLowering : public ConvertOpToLLVMPattern<GpuOp> {
     Value word3Const = createI32Constant(rewriter, loc, word3);
     resource = rewriter.create<LLVM::InsertElementOp>(
         loc, llvm4xI32, resource, word3Const,
-        this->createIndexConstant(rewriter, loc, 3));
+        this->createIndexAttrConstant(rewriter, loc, this->getIndexType(), 3));
     args.push_back(resource);
 
     // Indexing (voffset)

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

@@ -67,9 +67,10 @@ class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
 protected:
   Value getNumElements(ConversionPatternRewriter &rewriter, Location loc,
                        MemRefType type, MemRefDescriptor desc) const {
+    Type indexType = ConvertToLLVMPattern::getIndexType();
     return type.hasStaticShape()
-               ? ConvertToLLVMPattern::createIndexConstant(
-                     rewriter, loc, type.getNumElements())
+               ? ConvertToLLVMPattern::createIndexAttrConstant(
+                     rewriter, loc, indexType, type.getNumElements())
                // For identity maps (verified by caller), the number of
                // elements is stride[0] * size[0].
                : rewriter.create<LLVM::MulOp>(loc,

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -60,11 +60,6 @@ Value ConvertToLLVMPattern::createIndexAttrConstant(OpBuilder &builder,
                                           builder.getIndexAttr(value));
 }
 
-Value ConvertToLLVMPattern::createIndexConstant(
-    ConversionPatternRewriter &builder, Location loc, uint64_t value) const {
-  return createIndexAttrConstant(builder, loc, getIndexType(), value);
-}
-
 Value ConvertToLLVMPattern::getStridedElementPtr(
     Location loc, MemRefType type, Value memRefDesc, ValueRange indices,
     ConversionPatternRewriter &rewriter) const {
@@ -79,13 +74,15 @@ Value ConvertToLLVMPattern::getStridedElementPtr(
   Value base =
       memRefDescriptor.bufferPtr(rewriter, loc, *getTypeConverter(), type);
 
+  Type indexType = getIndexType();
   Value index;
   for (int i = 0, e = indices.size(); i < e; ++i) {
     Value increment = indices[i];
     if (strides[i] != 1) { // Skip if stride is 1.
-      Value stride = ShapedType::isDynamic(strides[i])
-                         ? memRefDescriptor.stride(rewriter, loc, i)
-                         : createIndexConstant(rewriter, loc, strides[i]);
+      Value stride =
+          ShapedType::isDynamic(strides[i])
+              ? memRefDescriptor.stride(rewriter, loc, i)
+              : createIndexAttrConstant(rewriter, loc, indexType, strides[i]);
       increment = rewriter.create<LLVM::MulOp>(loc, increment, stride);
     }
     index =
@@ -130,15 +127,17 @@ void ConvertToLLVMPattern::getMemRefDescriptorSizes(
 
   sizes.reserve(memRefType.getRank());
   unsigned dynamicIndex = 0;
+  Type indexType = getIndexType();
   for (int64_t size : memRefType.getShape()) {
-    sizes.push_back(size == ShapedType::kDynamic
-                        ? dynamicSizes[dynamicIndex++]
-                        : createIndexConstant(rewriter, loc, size));
+    sizes.push_back(
+        size == ShapedType::kDynamic
+            ? dynamicSizes[dynamicIndex++]
+            : createIndexAttrConstant(rewriter, loc, indexType, size));
   }
 
   // Strides: iterate sizes in reverse order and multiply.
   int64_t stride = 1;
-  Value runningStride = createIndexConstant(rewriter, loc, 1);
+  Value runningStride = createIndexAttrConstant(rewriter, loc, indexType, 1);
   strides.resize(memRefType.getRank());
   for (auto i = memRefType.getRank(); i-- > 0;) {
     strides[i] = runningStride;
@@ -158,7 +157,7 @@ void ConvertToLLVMPattern::getMemRefDescriptorSizes(
       runningStride =
           rewriter.create<LLVM::MulOp>(loc, runningStride, sizes[i]);
     else
-      runningStride = createIndexConstant(rewriter, loc, stride);
+      runningStride = createIndexAttrConstant(rewriter, loc, indexType, stride);
   }
   if (sizeInBytes) {
     // Buffer size in bytes.
@@ -195,22 +194,25 @@ Value ConvertToLLVMPattern::getNumElements(
              static_cast<ssize_t>(dynamicSizes.size()) &&
          "dynamicSizes size doesn't match dynamic sizes count in memref shape");
 
+  Type indexType = getIndexType();
   Value numElements = memRefType.getRank() == 0
-                          ? createIndexConstant(rewriter, loc, 1)
+                          ? createIndexAttrConstant(rewriter, loc, indexType, 1)
                           : nullptr;
   unsigned dynamicIndex = 0;
 
   // Compute the total number of memref elements.
   for (int64_t staticSize : memRefType.getShape()) {
     if (numElements) {
-      Value size = staticSize == ShapedType::kDynamic
-                       ? dynamicSizes[dynamicIndex++]
-                       : createIndexConstant(rewriter, loc, staticSize);
+      Value size =
+          staticSize == ShapedType::kDynamic
+              ? dynamicSizes[dynamicIndex++]
+              : createIndexAttrConstant(rewriter, loc, indexType, staticSize);
       numElements = rewriter.create<LLVM::MulOp>(loc, numElements, size);
     } else {
-      numElements = staticSize == ShapedType::kDynamic
-                        ? dynamicSizes[dynamicIndex++]
-                        : createIndexConstant(rewriter, loc, staticSize);
+      numElements =
+          staticSize == ShapedType::kDynamic
+              ? dynamicSizes[dynamicIndex++]
+              : createIndexAttrConstant(rewriter, loc, indexType, staticSize);
     }
   }
   return numElements;
@@ -231,8 +233,9 @@ MemRefDescriptor ConvertToLLVMPattern::createMemRefDescriptor(
   memRefDescriptor.setAlignedPtr(rewriter, loc, alignedPtr);
 
   // Field 3: Offset in aligned pointer.
-  memRefDescriptor.setOffset(rewriter, loc,
-                             createIndexConstant(rewriter, loc, 0));
+  Type indexType = getIndexType();
+  memRefDescriptor.setOffset(
+      rewriter, loc, createIndexAttrConstant(rewriter, loc, indexType, 0));
 
   // Fields 4: Sizes.
   for (const auto &en : llvm::enumerate(sizes))

mlir/lib/Conversion/MemRefToLLVM/AllocLikeConversion.cpp

@@ -138,7 +138,8 @@ bool AllocationOpLLVMLowering::isMemRefSizeMultipleOf(
 Value AllocationOpLLVMLowering::allocateBufferAutoAlign(
     ConversionPatternRewriter &rewriter, Location loc, Value sizeBytes,
     Operation *op, const DataLayout *defaultLayout, int64_t alignment) const {
-  Value allocAlignment = createIndexConstant(rewriter, loc, alignment);
+  Value allocAlignment =
+      createIndexAttrConstant(rewriter, loc, getIndexType(), alignment);
 
   MemRefType memRefType = getMemRefResultType(op);
   // Function aligned_alloc requires size to be a multiple of alignment; we pad