[mlir] Fix bitwidth of memref-to-llvm constant.

[gysit]

Summary

One constant generated in MemRefToLLVM had a hardcoded bitwidth of
64 bits. The fix uses the typeConverter to create a constant that
matches the bitwidth of the provided by the data layout. The issue was
detected in an attempt to add a verifier to the LLVM ICmp operation that
checks that the types of the compared arguments match.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

for (auto &en : llvm::enumerate(reassociation)) {
     }
   }
 }
 
 static void fillInStridesForCollapsedMemDescriptor(
     ConversionPatternRewriter &rewriter, Location loc, Operation *op,
     TypeConverter *typeConverter, MemRefType srcType, MemRefDescriptor &srcDesc,
     MemRefDescriptor &dstDesc, ArrayRef<ReassociationIndices> reassociation) {
+  auto llvmIndexType = typeConverter->convertType(rewriter.getIndexType());
   // See comments for computeCollapsedLayoutMap for details on how the strides
   // are calculated.
   auto srcShape = srcType.getShape();
   for (auto &en : llvm::enumerate(reassociation)) {
     rewriter.setInsertionPoint(op);
     auto dstIndex = en.index();
     ArrayRef<int64_t> ref = llvm::makeArrayRef(en.value());
     while (srcShape[ref.back()] == 1 && ref.size() > 1)
if (!ShapedType::isDynamic(srcShape[ref.back()]) || ref.size() == 1) {
         if (srcShape[srcIndex] == 1 && srcIndex != ref.front())
           continue;
         rewriter.setInsertionPointToEnd(curEntryBlock);
         Value srcStride = srcDesc.stride(rewriter, loc, srcIndex);
         if (srcIndex == ref.front()) {
           rewriter.create<LLVM::BrOp>(loc, srcStride, continueBlock);
           break;
         }
-        Value one = rewriter.create<LLVM::ConstantOp>(
-            loc, rewriter.getI64Type(), rewriter.getI32IntegerAttr(1));
+        Value one = rewriter.create<LLVM::ConstantOp>(loc, llvmIndexType,
+                                                      rewriter.getIndexAttr(1));
         Value predNeOne = rewriter.create<LLVM::ICmpOp>(
             loc, LLVM::ICmpPredicate::ne, srcDesc.size(rewriter, loc, srcIndex),
             one);
         {
           OpBuilder::InsertionGuard guard(rewriter);
           nextEntryBlock = rewriter.createBlock(
               initBlock->getParent(), Region::iterator(continueBlock), {});
         }

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 //       CHECK:      llvm.extractvalue %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:      llvm.mul %{{.*}}, %{{.*}}  : i64
 //       CHECK:      llvm.mlir.constant(4 : index) : i64
 //       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 //       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 //       CHECK:      llvm.extractvalue %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:      llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 //       CHECK:      llvm.extractvalue %{{.*}}[4, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
-//       CHECK:      llvm.mlir.constant(1 : i32) : i64
+//       CHECK:      llvm.mlir.constant(1 : index) : i64
 //       CHECK:      llvm.extractvalue %{{.*}}[3, 2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:      llvm.icmp "ne" %{{.*}}, %{{.*}} : i64
 //       CHECK:      llvm.cond_br %{{.*}}, ^bb2(%{{.*}} : i64), ^bb1
 //       CHECK:      ^bb1:
 //       CHECK:      llvm.extractvalue %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<3 x i64>, array<3 x i64>)>
 //       CHECK:      llvm.br ^bb2(%{{.*}} : i64)
 //       CHECK:      ^bb2(%[[STRIDE:.*]]: i64):
 //       CHECK:      llvm.insertvalue %[[STRIDE]], %{{.*}}[4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK32-LABEL: func @collapse_shape_dynamic_with_non_identity_layout(
+//       CHECK32:      llvm.mlir.constant(1 : index) : i32
+//       CHECK32:      llvm.mlir.constant(4 : index) : i32
+//       CHECK32:      llvm.mlir.constant(1 : index) : i32
 
 // -----
 
 func.func @expand_shape_static(%arg0: memref<3x4x5xf32>) -> memref<1x3x4x1x5xf32> {
   // Reshapes that expand a contiguous tensor with some 1's.
   %0 = memref.expand_shape %arg0 [[0, 1], [2], [3, 4]]
       : memref<3x4x5xf32> into memref<1x3x4x1x5xf32>
   return %0 : memref<1x3x4x1x5xf32>
func.func @memref_copy_unranked() {
   return
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_aligned_pointer_as_index
 func.func @extract_aligned_pointer_as_index(%m: memref<?xf32>) -> index {
   %0 = memref.extract_aligned_pointer_as_index %m: memref<?xf32> -> index
-  // CHECK: %[[E:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)> 
+  // CHECK: %[[E:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)>
   // CHECK: %[[I64:.*]] = llvm.ptrtoint %[[E]] : !llvm.ptr<f32> to i64
   // CHECK: %[[R:.*]] = builtin.unrealized_conversion_cast %[[I64]] : i64 to index
 
   // CHECK: return %[[R:.*]] : index
   return %0: index
 }