[mlir][memref] WIP - Reorganize conversions involving memref in the presence of backend-specific type conversions.

[nicolasvasilache]

Summary

This revision is a work in progress that tries to entangle the notion of backend-specific type conversion.
In particular it demonstrates that a notional sink pass (or generalized pattern application) helps
avoid footguns that inevitably arise from pass-based thinking that requires keeping all passes in sync with
the same type conversion decision.

[mehdi_amini]

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

51

This API should not be used for addressing then because it is decoupled from the memory space?

Won't this lead to unrealized_cast because of different index size between arithmetic and indexing?

60

It's a bit weird that the indexing of memref would differ from this API which is more generic and about address spaces.

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

58

This is just gonna do isSharedMemoryAddressSpace(type.getMemorySpace()); does this indirection pulls its weight here? Also shall it be defined inline?

[guraypp]

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

62

We could add hasConstantMemoryAddressSpace or hasLocalMemoryAddressSpace.

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

211

We should use 32b for the local and constant memory spaces as well. Here is the list of address spaces (I don't remember what is 2):

generic = 0
global = 1
shared = 3
constant = 4
local = 5

As far as I can see, MLIR does set shared address space, but does not set local or constant address spaces.

[Hardcode84]

Does all this complexity really belongs to llvm conversion pass?
Maybe add a separate pass which decomposes memrefs index calculations into explicit ops (possibly inserting appropriate trunc casts depending on memory space) and then result is just straightforwardly converted to llvm.
As as bonus, you will be able to additional canonicalizations/transformations between this new pass and llvm conversion.

mlir/include/mlir/Conversion/LLVMCommon/LoweringOptions.h

 //===- LoweringOptions.h - Common config for lowering to LLVM ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Provides a configuration shared by several conversions targeting the LLVM
 // dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_CONVERSION_LLVMCOMMON_LOWERINGOPTIONS_H
 #define MLIR_CONVERSION_LLVMCOMMON_LOWERINGOPTIONS_H
 
+#include "mlir/IR/BuiltinTypes.h"
 #include "llvm/IR/DataLayout.h"
 
 namespace mlir {
 
 class DataLayout;
 class MLIRContext;
 
 /// Value to pass as bitwidth for the index type when the converter is expected
void overrideIndexBitwidth(unsigned bitwidth) {
     assert(bitwidth != kDeriveIndexBitwidthFromDataLayout &&
            "can only override to a concrete bitwidth");
     indexBitwidth = bitwidth;
   }
 
   /// Get the index bitwidth.
   unsigned getIndexBitwidth() const { return indexBitwidth; }
 
+  /// Hook to customize the conversion of MemRefType to LLVMType.
+  llvm::function_ref<Type(MemRefType)> memrefIndexTypeConverter = nullptr;
+
 private:
   unsigned indexBitwidth;
 };
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_LLVMCOMMON_LOWERINGOPTIONS_H

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

 protected:
   /// Returns the LLVM dialect.
   LLVM::LLVMDialect &getDialect() const;
 
   LLVMTypeConverter *getTypeConverter() const;
 
   /// Gets the MLIR type wrapping the LLVM integer type whose bit width is
   /// defined by the used type converter.
   Type getIndexType() const;
 
+  /// Gets the MLIR type wrapping the LLVM integer type whose bit width is
+  /// defined by the used type converter and matching the index type needed for
+  /// MemRefType `t`.
+  Type getIndexTypeMatchingMemRef(MemRefType t) const;
+
   /// Gets the MLIR type wrapping the LLVM integer type whose bit width
   /// corresponds to that of a LLVM pointer type.
   Type getIntPtrType(unsigned addressSpace = 0) const;
 
   /// Gets the MLIR type wrapping the LLVM void type.
   Type getVoidType() const;
 
   /// Get the MLIR type wrapping the LLVM i8* type.
   Type getVoidPtrType() const;
 
   /// Create a constant Op producing a value of `resultType` from an index-typed

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

public:
   const DataLayoutAnalysis *getDataLayoutAnalysis() const {
     return dataLayoutAnalysis;
   }
 
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   Type getIndexType();
 
+  /// Gets the LLVM representation of the index type that matches the MemRefType
+  /// `t`. The returned type is an integer type with the size configured for
+  /// this type converter.
+  Type getIndexTypeMatchingMemRef(MemRefType t);
+
   /// Returns true if using opaque pointers was enabled in the lowering options.
   bool useOpaquePointers() const { return getOptions().useOpaquePointers; }
 
   /// Creates an LLVM pointer type with the given element type and address
   /// space.
   /// This function is meant to be used in code supporting both typed and opaque
   /// pointers, as it will create an opaque pointer with the given address space
   /// if opaque pointers are enabled in the lowering options.

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

protected:
 
   /// Computes the alignment for the given memory allocation op.
   template <typename OpType>
   Value getAlignment(ConversionPatternRewriter &rewriter, Location loc,
                      OpType op) const {
     MemRefType memRefType = op.getType();
     Value alignment;
     if (auto alignmentAttr = op.getAlignment()) {
-      Type indexType = getIndexType();
+      Type indexType =
+          ConvertToLLVMPattern::getIndexTypeMatchingMemRef(memRefType);
       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
       // biggest scalar on a target HW. For non-scalars, use the natural
       // alignment of the LLVM type given by the LLVM DataLayout.
       alignment = getSizeInBytes(loc, memRefType.getElementType(), rewriter);

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

let extraClassDeclaration = [{
 
     /// Returns the numeric value used to identify the workgroup memory address
     /// space.
     static AddressSpace getWorkgroupAddressSpace() { return AddressSpace::Workgroup; }
 
     /// Returns the numeric value used to identify the private memory address
     /// space.
     static AddressSpace getPrivateAddressSpace() { return AddressSpace::Private; }
+
+    /// Return true if the given MemRefType has an address space that is a 
+    /// gpu::AddressSpaceAttr attribute with value 'workgroup`.
+    static bool hasSharedMemoryAddressSpace(MemRefType type);
+
+    /// Return true if the given Attribute has matches is a gpu::AddressSpaceAttr
+    /// attribute with value 'workgroup`.
+    static bool isSharedMemoryAddressSpace(Attribute type);
   }];
 
   let dependentDialects = ["arith::ArithDialect"];
   let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
   let usePropertiesForAttributes = 1;
 }
 

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

 class ConvertOpToGpuRuntimeCallPattern : public ConvertOpToLLVMPattern<OpTy> {
 public:
   explicit ConvertOpToGpuRuntimeCallPattern(LLVMTypeConverter &typeConverter)
       : ConvertOpToLLVMPattern<OpTy>(typeConverter) {}
 
 protected:
   Value getNumElements(ConversionPatternRewriter &rewriter, Location loc,
                        MemRefType type, MemRefDescriptor desc) const {
-    Type indexType = ConvertToLLVMPattern::getIndexType();
+    Type indexType = ConvertToLLVMPattern::getIndexTypeMatchingMemRef(type);
     return type.hasStaticShape()
                ? 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,
                                               desc.stride(rewriter, loc, 0),
                                               desc.size(rewriter, loc, 0));

 private:
   LogicalResult
   matchAndRewrite(gpu::SDDMMOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
+static IntegerType getIndexTypeForMemRef(MemRefType t) {
+  int64_t numBits = gpu::GPUDialect::hasSharedMemoryAddressSpace(t) ? 32 : 64;
+  return IntegerType::get(t.getContext(), numBits);
+}
+
 void GpuToLLVMConversionPass::runOnOperation() {
   LowerToLLVMOptions options(&getContext());
   options.useOpaquePointers = useOpaquePointers;
   options.useBarePtrCallConv = hostBarePtrCallConv;
+  options.memrefIndexTypeConverter = getIndexTypeForMemRef;
 
   LLVMTypeConverter converter(&getContext(), options);
   RewritePatternSet patterns(&getContext());
   LLVMConversionTarget target(getContext());
 
   target.addIllegalDialect<gpu::GPUDialect>();
 
   mlir::arith::populateArithToLLVMConversionPatterns(converter, patterns);

mlir/lib/Conversion/GPUToNVVM/CMakeLists.txt

add_mlir_conversion_library(MLIRGPUToNVVMTransforms
   LINK_LIBS PUBLIC
   MLIRArithToLLVM
   MLIRFuncToLLVM
   MLIRGPUDialect
   MLIRGPUToGPURuntimeTransforms
   MLIRLLVMCommonConversion
   MLIRLLVMDialect
   MLIRMemRefToLLVM
+  MLIRNVGPUDialect
   MLIRNVVMDialect
   MLIRPass
   MLIRTransformUtils
+  MLIRVectorToLLVM
   )

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 
 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
+#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 #include "../GPUCommon/GPUOpsLowering.h"
 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
 #include "../GPUCommon/OpToFuncCallLowering.h"
 #include <optional>
 
matchAndRewrite(gpu::LaneIdOp op, gpu::LaneIdOp::Adaptor adaptor,
     rewriter.replaceOp(op, {newOp});
     return success();
   }
 };
 
 /// Import the GPU Ops to NVVM Patterns.
 #include "GPUToNVVM.cpp.inc"
 
+static IntegerType getIndexTypeForMemRef(MemRefType t) {
+  int64_t numBits = (gpu::GPUDialect::hasSharedMemoryAddressSpace(t) ||
+                     nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(t))
+                        ? 32
+                        : 64;
+  return IntegerType::get(t.getContext(), numBits);
+}
+
 /// A pass that replaces all occurrences of GPU device operations with their
 /// corresponding NVVM equivalent.
 ///
 /// This pass only handles device code and is not meant to be run on GPU host
 /// code.
 struct LowerGpuOpsToNVVMOpsPass
     : public impl::ConvertGpuOpsToNVVMOpsBase<LowerGpuOpsToNVVMOpsPass> {
   LowerGpuOpsToNVVMOpsPass() = default;
void runOnOperation() override {
     // Customize the bitwidth used for the device side index computations.
     LowerToLLVMOptions options(
         m.getContext(),
         DataLayout(cast<DataLayoutOpInterface>(m.getOperation())));
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
     options.useOpaquePointers = useOpaquePointers;
     options.useBarePtrCallConv = useBarePtrCallConv;
+    options.memrefIndexTypeConverter = getIndexTypeForMemRef;
 
     // Apply in-dialect lowering. In-dialect lowering will replace
     // ops which need to be lowered further, which is not supported by a
     // single conversion pass.
     {
       RewritePatternSet patterns(m.getContext());
       populateGpuRewritePatterns(patterns);
       if (failed(applyPatternsAndFoldGreedily(m, std::move(patterns))))
void runOnOperation() override {
     // Lowering for MMAMatrixType.
     converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
       return convertMMAToLLVMType(type);
     });
     RewritePatternSet llvmPatterns(m.getContext());
 
     arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);
+    populateVectorToLLVMConversionPatterns(converter, llvmPatterns);
     populateFuncToLLVMConversionPatterns(converter, llvmPatterns);
     populateFinalizeMemRefToLLVMConversionPatterns(converter, llvmPatterns);
     populateGpuToNVVMConversionPatterns(converter, llvmPatterns);
     populateGpuWMMAToNVVMConversionPatterns(converter, llvmPatterns);
     if (this->hasRedux)
       populateGpuSubgroupReduceOpLoweringPattern(converter, llvmPatterns);
     LLVMConversionTarget target(getContext());
     configureGpuToNVVMConversionLegality(target);

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

 #include "mlir/IR/BuiltinAttributes.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ConvertToLLVMPattern
 //===----------------------------------------------------------------------===//
 
+static Value convertToDesiredIndexType(OpBuilder &b, Location loc, Value src,
+                                       Type desiredIndexType) {
+  assert(src.getType().isIntOrIndex() && !src.getType().isIndex() &&
+         "expected int type");
+  assert(desiredIndexType.isIntOrIndex() && !desiredIndexType.isIndex() &&
+         "expected int type");
+  if (src.getType() == desiredIndexType)
+    return src;
+  if (src.getType().getIntOrFloatBitWidth() <
+      desiredIndexType.getIntOrFloatBitWidth())
+    return b.create<LLVM::SExtOp>(loc, desiredIndexType, src);
+  return b.create<LLVM::TruncOp>(loc, desiredIndexType, src);
+}
+
 ConvertToLLVMPattern::ConvertToLLVMPattern(StringRef rootOpName,
                                            MLIRContext *context,
                                            LLVMTypeConverter &typeConverter,
                                            PatternBenefit benefit)
     : ConversionPattern(typeConverter, rootOpName, benefit, context) {}
 
 LLVMTypeConverter *ConvertToLLVMPattern::getTypeConverter() const {
   return static_cast<LLVMTypeConverter *>(
       ConversionPattern::getTypeConverter());
 }
 
 LLVM::LLVMDialect &ConvertToLLVMPattern::getDialect() const {
   return *getTypeConverter()->getDialect();
 }
 
 Type ConvertToLLVMPattern::getIndexType() const {
   return getTypeConverter()->getIndexType();
 }
 
+Type ConvertToLLVMPattern::getIndexTypeMatchingMemRef(MemRefType t) const {
+  return getTypeConverter()->getIndexTypeMatchingMemRef(t);
+}
+
 Type ConvertToLLVMPattern::getIntPtrType(unsigned addressSpace) const {
   return IntegerType::get(&getTypeConverter()->getContext(),
                           getTypeConverter()->getPointerBitwidth(addressSpace));
 }
 
 Type ConvertToLLVMPattern::getVoidType() const {
   return LLVM::LLVMVoidType::get(&getTypeConverter()->getContext());
 }
Value ConvertToLLVMPattern::getStridedElementPtr(
   MemRefDescriptor memRefDescriptor(memRefDesc);
   // Use a canonical representation of the start address so that later
   // optimizations have a longer sequence of instructions to CSE.
   // If we don't do that we would sprinkle the memref.offset in various
   // position of the different address computations.
   Value base =
       memRefDescriptor.bufferPtr(rewriter, loc, *getTypeConverter(), type);
 
-  Type indexType = getIndexType();
+  Type indexType = getIndexTypeMatchingMemRef(type);
   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)
               : createIndexAttrConstant(rewriter, loc, indexType, strides[i]);
+      increment =
+          convertToDesiredIndexType(rewriter, loc, increment, indexType);
       increment = rewriter.create<LLVM::MulOp>(loc, increment, stride);
     }
+    increment = convertToDesiredIndexType(rewriter, loc, increment, indexType);
     index =
         index ? rewriter.create<LLVM::AddOp>(loc, index, increment) : increment;
   }
 
   Type elementPtrType = memRefDescriptor.getElementPtrType();
   return index ? rewriter.create<LLVM::GEPOp>(
                      loc, elementPtrType,
                      getTypeConverter()->convertType(type.getElementType()),
void ConvertToLLVMPattern::getMemRefDescriptorSizes(
   assert(isConvertibleAndHasIdentityMaps(memRefType) &&
          "layout maps must have been normalized away");
   assert(count(memRefType.getShape(), ShapedType::kDynamic) ==
              static_cast<ssize_t>(dynamicSizes.size()) &&
          "dynamicSizes size doesn't match dynamic sizes count in memref shape");
 
   sizes.reserve(memRefType.getRank());
   unsigned dynamicIndex = 0;
-  Type indexType = getIndexType();
+  Type indexType = getIndexTypeMatchingMemRef(memRefType);
   for (int64_t size : memRefType.getShape()) {
     sizes.push_back(
         size == ShapedType::kDynamic
             ? dynamicSizes[dynamicIndex++]
             : createIndexAttrConstant(rewriter, loc, indexType, size));
   }
 
   // Strides: iterate sizes in reverse order and multiply.

 
 Value ConvertToLLVMPattern::getNumElements(
     Location loc, MemRefType memRefType, ValueRange dynamicSizes,
     ConversionPatternRewriter &rewriter) const {
   assert(count(memRefType.getShape(), ShapedType::kDynamic) ==
              static_cast<ssize_t>(dynamicSizes.size()) &&
          "dynamicSizes size doesn't match dynamic sizes count in memref shape");
 
-  Type indexType = getIndexType();
+  Type indexType = getIndexTypeMatchingMemRef(memRefType);
   Value numElements = memRefType.getRank() == 0
                           ? createIndexAttrConstant(rewriter, loc, indexType, 1)
                           : nullptr;
   unsigned dynamicIndex = 0;
 
   // Compute the total number of memref elements.
   for (int64_t staticSize : memRefType.getShape()) {
     if (numElements) {
MemRefDescriptor ConvertToLLVMPattern::createMemRefDescriptor(
 
   // Field 1: Allocated pointer, used for malloc/free.
   memRefDescriptor.setAllocatedPtr(rewriter, loc, allocatedPtr);
 
   // Field 2: Actual aligned pointer to payload.
   memRefDescriptor.setAlignedPtr(rewriter, loc, alignedPtr);
 
   // Field 3: Offset in aligned pointer.
-  Type indexType = getIndexType();
+  Type indexType = getIndexTypeMatchingMemRef(memRefType);
   memRefDescriptor.setOffset(
       rewriter, loc, createIndexAttrConstant(rewriter, loc, indexType, 0));
 
   // Fields 4: Sizes.
   for (const auto &en : llvm::enumerate(sizes))
     memRefDescriptor.setSize(rewriter, loc, en.index(), en.value());
 
   // Field 5: Strides.

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

 MLIRContext &LLVMTypeConverter::getContext() {
   return *getDialect()->getContext();
 }
 
 Type LLVMTypeConverter::getIndexType() {
   return IntegerType::get(&getContext(), getIndexTypeBitwidth());
 }
 
+Type LLVMTypeConverter::getIndexTypeMatchingMemRef(MemRefType t) {
+  return options.memrefIndexTypeConverter ? options.memrefIndexTypeConverter(t)
+                                          : getIndexType();
+}
+
 LLVM::LLVMPointerType
 LLVMTypeConverter::getPointerType(Type elementType, unsigned int addressSpace) {
   if (useOpaquePointers())
     return LLVM::LLVMPointerType::get(&getContext(), addressSpace);
   return LLVM::LLVMPointerType::get(elementType, addressSpace);
 }
 
 unsigned LLVMTypeConverter::getPointerBitwidth(unsigned addressSpace) {
emitError(UnknownLoc::get(type.getContext()),
               "conversion of memref memory space ")
         << type.getMemorySpace()
         << " to integer address space "
            "failed. Consider adding memory space conversions.";
     return {};
   }
   auto ptrTy = getPointerType(elementType, *addressSpace);
 
-  auto indexTy = getIndexType();
+  Type indexTy = getIndexTypeMatchingMemRef(type);
 
   SmallVector<Type, 5> results = {ptrTy, ptrTy, indexTy};
   auto rank = type.getRank();
   if (rank == 0)
     return results;
 
   if (unpackAggregates)
     results.insert(results.end(), 2 * rank, indexTy);
   else
     results.insert(results.end(), 2, LLVM::LLVMArrayType::get(indexTy, rank));
   return results;
 }
 
 unsigned LLVMTypeConverter::getMemRefDescriptorSize(MemRefType type,
                                                     const DataLayout &layout) {
   // Compute the descriptor size given that of its components indicated above.
   unsigned space = *getMemRefAddressSpace(type);
   return 2 * llvm::divideCeil(getPointerBitwidth(space), 8) +
-         (1 + 2 * type.getRank()) * layout.getTypeSize(getIndexType());
+         (1 + 2 * type.getRank()) *
+             layout.getTypeSize(getIndexTypeMatchingMemRef(type));
 }
 
 /// Converts MemRefType to LLVMType. A MemRefType is converted to a struct that
 /// packs the descriptor fields as defined by `getMemRefDescriptorFields`.
 Type LLVMTypeConverter::convertMemRefType(MemRefType type) {
   // When converting a MemRefType to a struct with descriptor fields, do not
   // unpack the `sizes` and `strides` arrays.
   SmallVector<Type, 5> types =

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct ReallocOpLoweringBase : public AllocationOpLLVMLowering {
   LogicalResult matchAndRewrite(memref::ReallocOp op, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
     OpBuilder::InsertionGuard guard(rewriter);
     Location loc = op.getLoc();
 
     auto computeNumElements =
         [&](MemRefType type, function_ref<Value()> getDynamicSize) -> Value {
       // Compute number of elements.
-      Type indexType = ConvertToLLVMPattern::getIndexType();
+      Type indexType = ConvertToLLVMPattern::getIndexTypeMatchingMemRef(type);
       Value numElements =
           type.isDynamicDim(0)
               ? getDynamicSize()
               : createIndexAttrConstant(rewriter, loc, indexType,
                                         type.getDimSize(0));
       if (numElements.getType() != indexType)
         numElements = typeConverter->materializeTargetConversion(
             rewriter, loc, indexType, numElements);
Type indexPtrTy = getTypeConverter()->getPointerType(
         getTypeConverter()->getIndexType(), addressSpace);
     Value offsetPtr = rewriter.create<LLVM::GEPOp>(
         loc, indexPtrTy, elementType, scalarMemRefDescPtr,
         ArrayRef<LLVM::GEPArg>{0, 2});
 
     // The size value that we have to extract can be obtained using GEPop with
     // `dimOp.index() + 1` index argument.
     Value idxPlusOne = rewriter.create<LLVM::AddOp>(
-        loc, createIndexAttrConstant(rewriter, loc, getIndexType(), 1),
+        loc,
+        createIndexAttrConstant(rewriter, loc, adaptor.getIndex().getType(), 1),
         adaptor.getIndex());
     Value sizePtr = rewriter.create<LLVM::GEPOp>(
         loc, indexPtrTy, getTypeConverter()->getIndexType(), offsetPtr,
         idxPlusOne);
     return rewriter
         .create<LLVM::LoadOp>(loc, getTypeConverter()->getIndexType(), sizePtr)
         .getResult();
   }
private:
 
   Value extractSizeOfRankedMemRef(Type operandType, memref::DimOp dimOp,
                                   OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
     Location loc = dimOp.getLoc();
 
     // Take advantage if index is constant.
     MemRefType memRefType = cast<MemRefType>(operandType);
-    Type indexType = getIndexType();
+    Type indexType = getIndexTypeMatchingMemRef(memRefType);
     if (std::optional<int64_t> index = getConstantDimIndex(dimOp)) {
       int64_t i = *index;
       if (i >= 0 && i < memRefType.getRank()) {
         if (memRefType.isDynamicDim(i)) {
           // extract dynamic size from the memref descriptor.
           MemRefDescriptor descriptor(adaptor.getSource());
           return descriptor.size(rewriter, loc, i);
         }
if (shapeMemRefType.hasStaticShape()) {
         return rewriter.notifyMatchFailure(reshapeOp,
                                            "dynamic offset is unsupported");
 
       desc.setConstantOffset(rewriter, loc, offset);
 
       assert(targetMemRefType.getLayout().isIdentity() &&
              "Identity layout map is a precondition of a valid reshape op");
 
-      Type indexType = getIndexType();
+      Type indexType = getIndexTypeMatchingMemRef(targetMemRefType);
       Value stride = nullptr;
       int64_t targetRank = targetMemRefType.getRank();
       for (auto i : llvm::reverse(llvm::seq<int64_t>(0, targetRank))) {
         if (!ShapedType::isDynamic(strides[i])) {
           // If the stride for this dimension is dynamic, then use the product
           // of the sizes of the inner dimensions.
           stride =
               createIndexAttrConstant(rewriter, loc, indexType, strides[i]);
convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
 
     // Use the offset pointer as base for further addressing. Copy over the new
     // shape and compute strides. For this, we create a loop from rank-1 to 0.
     Value targetSizesBase = UnrankedMemRefDescriptor::sizeBasePtr(
         rewriter, loc, *getTypeConverter(), underlyingDescPtr, elementPtrType);
     Value targetStridesBase = UnrankedMemRefDescriptor::strideBasePtr(
         rewriter, loc, *getTypeConverter(), targetSizesBase, resultRank);
     Value shapeOperandPtr = shapeDesc.alignedPtr(rewriter, loc);
-    Value oneIndex = createIndexAttrConstant(rewriter, loc, getIndexType(), 1);
+    Value oneIndex =
+        createIndexAttrConstant(rewriter, loc, resultRank.getType(), 1);
     Value resultRankMinusOne =
         rewriter.create<LLVM::SubOp>(loc, resultRank, oneIndex);
 
     Block *initBlock = rewriter.getInsertionBlock();
     Type indexType = getTypeConverter()->getIndexType();
     Block::iterator remainingOpsIt = std::next(rewriter.getInsertionPoint());
 
     Block *condBlock = rewriter.createBlock(initBlock->getParent(), {},
matchAndRewrite(memref::ViewOp viewOp, OpAdaptor adaptor,
     } else {
       bitcastPtr = rewriter.create<LLVM::BitcastOp>(
           loc, LLVM::LLVMPointerType::get(targetElementTy, sourceMemorySpace),
           alignedPtr);
     }
 
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
-    Type indexType = getIndexType();
+    auto indexType = targetMemRef.getIndexType();
     // Field 3: The offset in the resulting type must be 0. This is
     // because of the type change: an offset on srcType* may not be
     // expressible as an offset on dstType*.
     targetMemRef.setOffset(
         rewriter, loc,
         createIndexAttrConstant(rewriter, loc, indexType, offset));
 
     // Early exit for 0-D corner case.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 
 using namespace mlir;
 using namespace mlir::gpu;
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.cpp.inc"
 
+/// Return true if the given MemRefType has an address space that is a
+/// gpu::AddressSpaceAttr attribute with value 'workgroup`.
+bool gpu::GPUDialect::hasSharedMemoryAddressSpace(MemRefType type) {
+  return isSharedMemoryAddressSpace(type.getMemorySpace());
+}
+
+/// Return true if the given Attribute has matches is a gpu::AddressSpaceAttr
+/// attribute with value 'workgroup`.
+bool gpu::GPUDialect::isSharedMemoryAddressSpace(Attribute memorySpace) {
+  if (!memorySpace)
+    return false;
+  if (auto gpuAttr = llvm::dyn_cast<gpu::AddressSpaceAttr>(memorySpace))
+    return gpuAttr.getValue() == gpu::AddressSpace::Workgroup;
+  return false;
+}
+
 //===----------------------------------------------------------------------===//
 // GPU Device Mapping Attributes
 //===----------------------------------------------------------------------===//
 
 int64_t GPUBlockMappingAttr::getMappingId() const {
   return static_cast<int64_t>(getBlock());
 }
 

mlir/test/Conversion/GPUCommon/memory-attrbution.mlir

gpu.func @workgroup(%arg0: f32) workgroup(%arg1: memref<4xf32, #gpu.address_space<workgroup>>) {
     // NVVM: %[[raw:.*]] = llvm.getelementptr %[[addr]][0, 0]
     // NVVM-SAME: !llvm.ptr<3>
 
     // ROCDL: %[[addr:.*]] = llvm.mlir.addressof @[[$buffer]] : !llvm.ptr<3>
     // ROCDL: %[[raw:.*]] = llvm.getelementptr %[[addr]][0, 0]
     // ROCDL-SAME: !llvm.ptr<3>
 
     // Populate the memref descriptor.
-    // NVVM: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>
+    // NVVM: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i32, array<1 x i32>, array<1 x i32>)>
     // NVVM: %[[descr2:.*]] = llvm.insertvalue %[[raw]], %[[descr1]][0]
     // NVVM: %[[descr3:.*]] = llvm.insertvalue %[[raw]], %[[descr2]][1]
-    // NVVM: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
+    // NVVM: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i32
     // NVVM: %[[descr4:.*]] = llvm.insertvalue %[[c0]], %[[descr3]][2]
-    // NVVM: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i64
+    // NVVM: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i32
     // NVVM: %[[descr5:.*]] = llvm.insertvalue %[[c4]], %[[descr4]][3, 0]
-    // NVVM: %[[c1:.*]] = llvm.mlir.constant(1 : index) : i64
+    // NVVM: %[[c1:.*]] = llvm.mlir.constant(1 : index) : i32
     // NVVM: %[[descr6:.*]] = llvm.insertvalue %[[c1]], %[[descr5]][4, 0]
 
     // ROCDL: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>
     // ROCDL: %[[descr2:.*]] = llvm.insertvalue %[[raw]], %[[descr1]][0]
     // ROCDL: %[[descr3:.*]] = llvm.insertvalue %[[raw]], %[[descr2]][1]
     // ROCDL: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
     // ROCDL: %[[descr4:.*]] = llvm.insertvalue %[[c0]], %[[descr3]][2]
     // ROCDL: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i64
gpu.func @workgroup3d(%arg0: f32) workgroup(%arg1: memref<4x2x6xf32, #gpu.address_space<workgroup>>) {
     // NVVM: %[[raw:.*]] = llvm.getelementptr %[[addr]][0, 0]
     // NVVM-SAME: !llvm.ptr<3>
 
     // ROCDL: %[[addr:.*]] = llvm.mlir.addressof @[[$buffer]] : !llvm.ptr<3>
     // ROCDL: %[[raw:.*]] = llvm.getelementptr %[[addr]][0, 0]
     // ROCDL-SAME: !llvm.ptr<3>
 
     // Populate the memref descriptor.
-    // NVVM: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i64, array<3 x i64>, array<3 x i64>)>
+    // NVVM: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i32, array<3 x i32>, array<3 x i32>)>
     // NVVM: %[[descr2:.*]] = llvm.insertvalue %[[raw]], %[[descr1]][0]
     // NVVM: %[[descr3:.*]] = llvm.insertvalue %[[raw]], %[[descr2]][1]
-    // NVVM: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
+    // NVVM: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i32
     // NVVM: %[[descr4:.*]] = llvm.insertvalue %[[c0]], %[[descr3]][2]
-    // NVVM: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i64
+    // NVVM: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i32
     // NVVM: %[[descr5:.*]] = llvm.insertvalue %[[c4]], %[[descr4]][3, 0]
-    // NVVM: %[[c12:.*]] = llvm.mlir.constant(12 : index) : i64
+    // NVVM: %[[c12:.*]] = llvm.mlir.constant(12 : index) : i32
     // NVVM: %[[descr6:.*]] = llvm.insertvalue %[[c12]], %[[descr5]][4, 0]
-    // NVVM: %[[c2:.*]] = llvm.mlir.constant(2 : index) : i64
+    // NVVM: %[[c2:.*]] = llvm.mlir.constant(2 : index) : i32
     // NVVM: %[[descr7:.*]] = llvm.insertvalue %[[c2]], %[[descr6]][3, 1]
-    // NVVM: %[[c6:.*]] = llvm.mlir.constant(6 : index) : i64
+    // NVVM: %[[c6:.*]] = llvm.mlir.constant(6 : index) : i32
     // NVVM: %[[descr8:.*]] = llvm.insertvalue %[[c6]], %[[descr7]][4, 1]
-    // NVVM: %[[c6:.*]] = llvm.mlir.constant(6 : index) : i64
+    // NVVM: %[[c6:.*]] = llvm.mlir.constant(6 : index) : i32
     // NVVM: %[[descr9:.*]] = llvm.insertvalue %[[c6]], %[[descr8]][3, 2]
-    // NVVM: %[[c1:.*]] = llvm.mlir.constant(1 : index) : i64
+    // NVVM: %[[c1:.*]] = llvm.mlir.constant(1 : index) : i32
     // NVVM: %[[descr10:.*]] = llvm.insertvalue %[[c1]], %[[descr9]][4, 2]
 
     // ROCDL: %[[descr1:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<3>, ptr<3>, i64, array<3 x i64>, array<3 x i64>)>
     // ROCDL: %[[descr2:.*]] = llvm.insertvalue %[[raw]], %[[descr1]][0]
     // ROCDL: %[[descr3:.*]] = llvm.insertvalue %[[raw]], %[[descr2]][1]
     // ROCDL: %[[c0:.*]] = llvm.mlir.constant(0 : index) : i64
     // ROCDL: %[[descr4:.*]] = llvm.insertvalue %[[c0]], %[[descr3]][2]
     // ROCDL: %[[c4:.*]] = llvm.mlir.constant(4 : index) : i64

mlir/test/Conversion/GPUToNVVM/typed-pointers.mlir

 // RUN: mlir-opt --convert-gpu-to-nvvm="use-opaque-pointers=0" --split-input-file %s | FileCheck %s
 // RUN: mlir-opt --convert-gpu-to-nvvm="index-bitwidth=32 use-opaque-pointers=0" --split-input-file %s | FileCheck --check-prefix=CHECK32 %s
 
 gpu.module @test_module {
 
   // CHECK-LABEL: func @gpu_wmma_load_op() ->
   // CHECK-SAME: !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
   // CHECK32-LABEL: func @gpu_wmma_load_op() ->
   func.func @gpu_wmma_load_op() -> (!gpu.mma_matrix<16x16xf16, "AOp">) {
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     %i = arith.constant 16 : index
     %j = arith.constant 16 : index
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %j] {leadDimension = 32 : index, transpose} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
-    // CHECK:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i64
+    // CHECK:  %[[INX64:.*]] = llvm.mlir.constant(16 : index) : i64
     // CHECK: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
-    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f16, 3>, ptr<f16, 3>, i64, array<2 x i64>, array<2 x i64>)>
-    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i64
-    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i64
-    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX]]  : i64
-    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<f16, 3>, i64) -> !llvm.ptr<f16, 3>
+    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f16, 3>, ptr<f16, 3>, i32, array<2 x i32>, array<2 x i32>)>
+    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i32
+    // CHECK:  %[[INX:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i32
+    // CHECK:  %[[INX2:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX2]]  : i32
+    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<f16, 3>, i32) -> !llvm.ptr<f16, 3>
     // CHECK:  %[[LDM32:.*]] = llvm.mlir.constant(32 : index) : i32
     // CHECK:  %[[FRAG:.*]] = nvvm.wmma.load %[[ADDRESS]], %[[LDM32]]
     // CHECK-SAME: {eltype = #nvvm.mma_type<f16>, frag = #nvvm.mma_frag<a>, k = 16 : i32, layout = #nvvm.mma_layout<col>, m = 16 : i32, n = 16 : i32}  : (!llvm.ptr<f16, 3>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
     // CHECK:  llvm.return %[[FRAG]] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
 
     // CHECK32:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i32
     // CHECK32: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK32:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<f16, 3>, ptr<f16, 3>, i32, array<2 x i32>, array<2 x i32>)>

mlir/test/Conversion/GPUToNVVM/wmma-ops-to-nvvm.mlir

 // RUN: mlir-opt --convert-gpu-to-nvvm='use-opaque-pointers=1' --split-input-file %s | FileCheck %s
 // RUN: mlir-opt --convert-gpu-to-nvvm="index-bitwidth=32 use-opaque-pointers=1" --split-input-file %s | FileCheck --check-prefix=CHECK32 %s
 
 gpu.module @test_module {
 
   // CHECK-LABEL: func @gpu_wmma_load_op() ->
   // CHECK-SAME: !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
   // CHECK32-LABEL: func @gpu_wmma_load_op() ->
   func.func @gpu_wmma_load_op() -> (!gpu.mma_matrix<16x16xf16, "AOp">) {
     %wg = memref.alloca() {alignment = 32} : memref<32x32xf16, 3>
     %i = arith.constant 16 : index
     %j = arith.constant 16 : index
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %j] {leadDimension = 32 : index, transpose} : memref<32x32xf16, 3> -> !gpu.mma_matrix<16x16xf16, "AOp">
-    // CHECK:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i64
+    // CHECK:  %[[INX64:.*]] = llvm.mlir.constant(16 : index) : i64
     // CHECK: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
-    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>
-    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i64
-    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i64
-    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX]]  : i64
-    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f16
+    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i32, array<2 x i32>, array<2 x i32>)>
+    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i32
+    // CHECK:  %[[INX:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i32
+    // CHECK:  %[[INX2:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX2]]  : i32
+    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, f16
     // CHECK:  %[[LDM32:.*]] = llvm.mlir.constant(32 : index) : i32
     // CHECK:  %[[FRAG:.*]] = nvvm.wmma.load %[[ADDRESS]], %[[LDM32]]
     // CHECK-SAME: {eltype = #nvvm.mma_type<f16>, frag = #nvvm.mma_frag<a>, k = 16 : i32, layout = #nvvm.mma_layout<col>, m = 16 : i32, n = 16 : i32}  : (!llvm.ptr<3>) -> !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
     // CHECK:  llvm.return %[[FRAG]] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
 
     // CHECK32:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i32
     // CHECK32: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK32:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i32, array<2 x i32>, array<2 x i32>)>
gpu.module @test_module {
   // CHECK-LABEL: func @gpu_wmma_int8_load_op() ->
   // CHECK-SAME: !llvm.struct<(i32, i32)>
   // CHECK32-LABEL: func @gpu_wmma_int8_load_op() ->
   func.func @gpu_wmma_int8_load_op() -> (!gpu.mma_matrix<16x16xsi8, "AOp">) {
     %wg = memref.alloca() {alignment = 32} : memref<32x32xi8, 3>
     %i = arith.constant 16 : index
     %j = arith.constant 16 : index
     %0 = gpu.subgroup_mma_load_matrix %wg[%i, %j] {leadDimension = 32 : index, transpose} : memref<32x32xi8, 3> -> !gpu.mma_matrix<16x16xsi8, "AOp">
-    // CHECK:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i64
+    // CHECK:  %[[INX64:.*]] = llvm.mlir.constant(16 : index) : i64
     // CHECK: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
-    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>
-    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i64
-    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i64
-    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX]]  : i64
-    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, i8
+    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i32, array<2 x i32>, array<2 x i32>)>
+    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i32
+    // CHECK:  %[[INX:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]  : i32
+    // CHECK:  %[[INX2:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX2]]  : i32
+    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, i8
     // CHECK:  %[[LDM32:.*]] = llvm.mlir.constant(32 : index) : i32
     // CHECK:  %[[FRAG:.*]] = nvvm.wmma.load %[[ADDRESS]], %[[LDM32]]
     // CHECK-SAME: {eltype = #nvvm.mma_type<s8>, frag = #nvvm.mma_frag<a>, k = 16 : i32, layout = #nvvm.mma_layout<col>, m = 16 : i32, n = 16 : i32}  : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32)>
     // CHECK:  llvm.return %[[FRAG]] : !llvm.struct<(i32, i32)>
 
     // CHECK32:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i32
     // CHECK32: %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK32:  %[[BASE:.*]] = llvm.extractvalue %{{.*}}[1] : !llvm.struct<(ptr<3>, ptr<3>, i32, array<2 x i32>, array<2 x i32>)>
gpu.module @test_module {
   // CHECK-SAME: (%[[D:.*]]: !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>)
   // CHECK32-LABEL: func @gpu_wmma_store_op
   // CHECK32-SAME: (%[[D:.*]]: !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>)
   func.func @gpu_wmma_store_op(%arg0 : !gpu.mma_matrix<16x16xf16, "COp">) -> () {
     %sg = memref.alloca(){alignment = 32} : memref<32x32xf16, 3>
     %i = arith.constant 16 : index
     %j = arith.constant 16 : index
     gpu.subgroup_mma_store_matrix %arg0, %sg[%i,%j] {leadDimension= 32 : index, transpose} : !gpu.mma_matrix<16x16xf16, "COp">, memref<32x32xf16, 3>
-    // CHECK:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i64
+    // CHECK:  %[[INX64:.*]] = llvm.mlir.constant(16 : index) : i64
     // CHECK:  %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK:  %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK:  %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK:  %[[MEMREF:.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK:  %[[EL1:.*]] = llvm.extractvalue %[[D]][0] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
     // CHECK:  %[[EL2:.*]] = llvm.extractvalue %[[D]][1] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
     // CHECK:  %[[EL3:.*]] = llvm.extractvalue %[[D]][2] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
     // CHECK:  %[[EL4:.*]] = llvm.extractvalue %[[D]][3] : !llvm.struct<(vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>)>
-    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %[[MEMREF]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<2 x i64>, array<2 x i64>)>
-    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i64
-    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]   : i64
-    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX]]  : i64
-    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f16
+    // CHECK:  %[[BASE:.*]] = llvm.extractvalue %[[MEMREF]][1] : !llvm.struct<(ptr<3>, ptr<3>, i32, array<2 x i32>, array<2 x i32>)>
+    // CHECK:  %[[LDM:.*]] = llvm.mlir.constant(32 : index) : i32
+    // CHECK:  %[[INX:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LI:.*]] = llvm.mul %[[INX]], %[[LDM]]   : i32
+    // CHECK:  %[[INX2:.*]] = llvm.trunc %[[INX64]] : i64 to i32
+    // CHECK:  %[[LIJ:.*]] = llvm.add %[[LI]], %[[INX2]]  : i32
+    // CHECK:  %[[ADDRESS:.*]] = llvm.getelementptr %[[BASE]][%[[LIJ]]] : (!llvm.ptr<3>, i32) -> !llvm.ptr<3>, f16
     // CHECK:  %[[LDM32:.*]] = llvm.mlir.constant(32 : index) : i32
     // CHECK:  nvvm.wmma.store %[[ADDRESS]], %[[LDM32]], %[[EL1]], %[[EL2]], %[[EL3]], %[[EL4]]
     // CHECK-SAME: {eltype = #nvvm.mma_type<f16>, k = 16 : i32, layout = #nvvm.mma_layout<col>, m = 16 : i32, n = 16 : i32} : !llvm.ptr<3>, vector<2xf16>, vector<2xf16>, vector<2xf16>, vector<2xf16>
     // CHECK:  llvm.return
 
     // CHECK32:  %[[INX:.*]] = llvm.mlir.constant(16 : index) : i32
     // CHECK32:  %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]
     // CHECK32:  %{{.*}} = llvm.insertvalue %{{.*}}, %{{.*}}[{{.*}}, {{.*}}]

mlir/test/Integration/Dialect/Vector/GPU/CUDA/test-warp-distribute.mlir

 // Run the test cases without distributing ops to test default lowering. Run
 // everything on the same thread.
 // RUN: mlir-opt %s -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
-// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
+// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-arith-to-llvm \
 // RUN:  -gpu-kernel-outlining |\
 // RUN: mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin))' |\
 // RUN: mlir-opt -gpu-to-llvm -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_cuda_runtime \
 // RUN:   -shared-libs=%mlir_c_runner_utils \
 // RUN:   -shared-libs=%mlir_runner_utils | \
 // RUN: FileCheck %s
 
 // Run the same test cases with distribution and propagation.
 // RUN: mlir-opt %s  -test-vector-warp-distribute="hoist-uniform distribute-transfer-write" \
 // RUN:   -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
-// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
+// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-arith-to-llvm \
 // RUN:  -gpu-kernel-outlining |\
 // RUN: mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin))' |\
 // RUN: mlir-opt -gpu-to-llvm -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_cuda_runtime \
 // RUN:   -shared-libs=%mlir_c_runner_utils \
 // RUN:   -shared-libs=%mlir_runner_utils | \
 // RUN: FileCheck %s
 
 // RUN: mlir-opt %s  -test-vector-warp-distribute="hoist-uniform distribute-transfer-write propagate-distribution" \
 // RUN:   -test-vector-warp-distribute=rewrite-warp-ops-to-scf-if -canonicalize | \
-// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-vector-to-llvm -convert-arith-to-llvm \
+// RUN: mlir-opt -convert-scf-to-cf -convert-cf-to-llvm -convert-arith-to-llvm \
 // RUN:  -gpu-kernel-outlining |\
 // RUN: mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,reconcile-unrealized-casts,gpu-to-cubin))' |\
 // RUN: mlir-opt -gpu-to-llvm -reconcile-unrealized-casts |\
 // RUN: mlir-cpu-runner -e main -entry-point-result=void \
 // RUN:   -shared-libs=%mlir_cuda_runtime \
 // RUN:   -shared-libs=%mlir_c_runner_utils \
 // RUN:   -shared-libs=%mlir_runner_utils | \
 // RUN: FileCheck %s

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":GPUToNVVMGen",
         ":GPUTransforms",
         ":IR",
         ":LLVMCommonConversion",
         ":LLVMDialect",
         ":MathDialect",
         ":MemRefDialect",
         ":MemRefToLLVM",
+        ":NVGPUDialect",
         ":NVVMDialect",
         ":Pass",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 cc_library(