[mlir][NVGPU] WIP - Apply layering changes for e2e NVVM

[nicolasvasilache]

Summary

Better index -> i32/i64 handling is needed to support TMA operations.
This commit is a draft that makes some connections possible.

[nicolasvasilache]

@guraypp @mehdi_amini @kerrmudgeon FYI, with this WIP I am able to generate PTX for TMA descriptors + load based on @guraypp 's ongoing work

[nicolasvasilache]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

657

what is the proper way to replace the magic constant at the GPU -> LLVM conversion ?

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

206

what is the proper way to replace the magic constant at the GPU -> NVVM conversion ?

[guraypp]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

657

I put isSharedMemoryAddressSpace in nvgpu dialect, not sure you can use it here.

[nicolasvasilache]

Update after debugging with --mlir-print-ir-after-all and ensuring all intermediate IRs are valid.

[qcolombet]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

659

The size of the address space pointer is technically carried by the datalayout in LLVM.

For NVPTX in particular, this is controlled by nvptx-short-ptr (https://github.com/llvm/llvm-project/blob/0b17e9d2859acfec2cf757472f3822f6b5aad020/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp#L60).

So we may need to query the datalayout here (no idea if that's even feasible) instead of hardcoding this.

[guraypp]

mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

659

I previously highlighted this flag in another PR. Let's ensure its setting in MLIR, otherwise LLVM promotes 32-bit registers to 64-bit no matter what we do in MLIR.

Nevertheless, having 32-bit in MLIR, as in this work, offers advantages. For instance, when generating assembly or PTX directly, maintaining 32-bit is crucial.

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/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp

 private:
   LogicalResult
   matchAndRewrite(gpu::SDDMMOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
+static IntegerType getIndexTypeForMemRef(MemRefType type) {
+  if (type.getMemorySpaceAsInt() == 3)
+    // nvgpu::NVGPUDialect::kSharedMemoryAddressSpace)
+    return IntegerType::get(type.getContext(), 32);
+  return IntegerType::get(type.getContext(), 64);
+}
+
 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/LowerGpuOpsToNVVMOps.cpp

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 type) {
+  if (type.getMemorySpaceAsInt() == 3)
+    // nvgpu::NVGPUDialect::kSharedMemoryAddressSpace)
+    return IntegerType::get(type.getContext(), 32);
+  return IntegerType::get(type.getContext(), 64);
+}
+
 /// 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))))

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

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();
+  auto indexTy = options.memrefIndexTypeConverter
+                     ? options.memrefIndexTypeConverter(type)
+                     : getIndexType();
 
   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);

mlir/lib/Conversion/MemRefToLLVM/CMakeLists.txt

 add_mlir_conversion_library(MLIRMemRefToLLVM
   AllocLikeConversion.cpp
   MemRefToLLVM.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/MemRefToLLVM
 
   DEPENDS
   MLIRConversionPassIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRAnalysis
   MLIRDataLayoutInterfaces
+  MLIRIndexDialect
   MLIRLLVMCommonConversion
   MLIRMemRefDialect
   MLIRMemRefUtils
   MLIRLLVMDialect
   MLIRTransforms
   )

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

 //===- NVGPUToNVVM.cpp - NVGPU to NVVM dialect conversion -----------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
 
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/Index/IR/IndexDialect.h"
+#include "mlir/Dialect/Index/IR/IndexOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_CONVERTNVGPUTONVVMPASS
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 
-/// GPU has 32 bit registers, this function truncates values when larger width
-/// is not needed.
+/// GPU has 32 bit registers, this function truncates values when larger
+/// width is not needed.
 static Value truncToI32(ConversionPatternRewriter &rewriter, Location loc,
                         Value value) {
   Type type = value.getType();
+  if (llvm::isa<IndexType>(type))
+    return rewriter.create<index::CastSOp>(loc, rewriter.getI32Type(), value);
+
   assert(llvm::isa<IntegerType>(type) && "expected an integer Value");
   if (type.getIntOrFloatBitWidth() <= 32)
     return value;
-  return rewriter.create<LLVM::TruncOp>(loc, rewriter.getI32Type(), value);
+  // Avoid direct use of LVVM and instead roundtrip through index dialect which
+  // connects things properly.
+  Value index =
+      rewriter.create<index::CastSOp>(loc, rewriter.getIndexType(), value);
+  return rewriter.create<index::CastSOp>(loc, rewriter.getI32Type(), index);
 }
 
 /// Returns the type for the intrinsic given the vectorResultType of the
 /// `gpu.mma.sync` operation.
 static Type inferIntrinsicResultType(Type vectorResultType) {
   MLIRContext *ctx = vectorResultType.getContext();
   auto a = cast<LLVM::LLVMArrayType>(vectorResultType);
   auto f16x2Ty = LLVM::getFixedVectorType(Float16Type::get(ctx), 2);
static Value convertIntrinsicResult(Location loc, Type intrinsicResultType,
   Type f64Ty = rewriter.getF64Type();
   Type f16x2Ty = LLVM::getFixedVectorType(rewriter.getF16Type(), 2);
   Type i32x2Ty = LLVM::getFixedVectorType(i32Ty, 2);
   Type f64x2Ty = LLVM::getFixedVectorType(f64Ty, 2);
   Type f32x2Ty = LLVM::getFixedVectorType(f32Ty, 2);
   Type f32x1Ty = LLVM::getFixedVectorType(f32Ty, 1);
 
   auto makeConst = [&](int32_t index) -> Value {
-    return rewriter.create<LLVM::ConstantOp>(loc, IntegerType::get(ctx, 32),
-                                             rewriter.getI32IntegerAttr(index));
+    return rewriter.create<index::ConstantOp>(loc,
+                                              rewriter.getIndexAttr(index));
   };
 
   if (arrayType) {
     SmallVector<Value, 4> elements;
 
     // The intrinsic returns 32-bit wide elements in a form which can be
     // directly bitcasted and inserted into the result vector.
     if (arrayType.getElementType() == f16x2Ty ||
for (unsigned i = 0, e = arrayTy.getNumElements(); i < e; ++i) {
     VectorType innerArrayTy = dyn_cast<VectorType>(arrayTy.getElementType());
     if (innerArrayTy && (innerArrayTy.getElementType() == i32Ty ||
                          innerArrayTy.getElementType() == f64Ty ||
                          innerArrayTy.getElementType() == f32Ty)) {
       for (unsigned idx = 0, innerSize = innerArrayTy.getNumElements();
            idx < innerSize; idx++) {
         result.push_back(rewriter.create<LLVM::ExtractElementOp>(
             loc, toUse,
-            rewriter.create<LLVM::ConstantOp>(
-                loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(idx))));
+            rewriter.create<index::ConstantOp>(loc,
+                                               rewriter.getIndexAttr(idx))));
       }
       continue;
     }
     result.push_back(toUse);
   }
   return result;
 }
 
static Value getMbarrierPtr(ConversionPatternRewriter &rewriter,
                             TypedValue<nvgpu::MBarrierType> barrier,
                             Value barrierMemref) {
   MemRefType memrefType = createMBarrierMemrefType(rewriter, barrier.getType());
   MemRefDescriptor memRefDescriptor(barrierMemref);
   return memRefDescriptor.bufferPtr(rewriter, barrier.getLoc(), typeConverter,
                                     memrefType);
 }
 
-struct ConvertNVGPUToNVVMPass
-    : public impl::ConvertNVGPUToNVVMPassBase<ConvertNVGPUToNVVMPass> {
-  using Base::Base;
-
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry
-        .insert<memref::MemRefDialect, LLVM::LLVMDialect, NVVM::NVVMDialect>();
-  }
-
-  void runOnOperation() override {
-    LowerToLLVMOptions options(&getContext());
-    options.useOpaquePointers = useOpaquePointers;
-    RewritePatternSet patterns(&getContext());
-    LLVMTypeConverter converter(&getContext(), options);
-    IRRewriter rewriter(&getContext());
-    /// device-side async tokens cannot be materialized in nvvm. We just
-    /// convert them to a dummy i32 type in order to easily drop them during
-    /// conversion.
-    converter.addConversion([&](nvgpu::DeviceAsyncTokenType type) -> Type {
-      return converter.convertType(IntegerType::get(type.getContext(), 32));
-    });
-    converter.addConversion([&](nvgpu::MBarrierTokenType type) -> Type {
-      return converter.convertType(IntegerType::get(type.getContext(), 64));
-    });
-    converter.addConversion([&](nvgpu::MBarrierType type) -> Type {
-      return converter.convertType(createMBarrierMemrefType(rewriter, type));
-    });
-    converter.addConversion([&](nvgpu::TensorMapDescriptorType type) -> Type {
-      return converter.getPointerType(type.getTensor().getElementType());
-    });
-    populateNVGPUToNVVMConversionPatterns(converter, patterns);
-    LLVMConversionTarget target(getContext());
-    target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
-    target.addLegalDialect<::mlir::memref::MemRefDialect>();
-    target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
-    if (failed(applyPartialConversion(getOperation(), target,
-                                      std::move(patterns))))
-      signalPassFailure();
-  }
-};
-
 /// Returns the constraints for the sparse MMA inline assembly instruction.
 static std::string buildMmaSparseAsmConstraintString(unsigned matASize,
                                                      unsigned matBSize,
                                                      unsigned matCSize) {
   std::string str;
   llvm::raw_string_ostream ss(str);
   for (unsigned i = 0; i < matCSize; i++)
     ss << "=r,";
matchAndRewrite(nvgpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
     // memory) to fill DstElements number of elements in the destination
     // (shared memory).
     Value srcBytes = adaptor.getSrcElements();
     if (srcBytes) {
       // When the optional SrcElements argument is present, the source (global
       // memory) of CpAsyncOp is read only for SrcElements number of elements.
       // The rest of the DstElements in the destination (shared memory) are
       // filled with zeros.
-      Value c3I32 = rewriter.create<LLVM::ConstantOp>(
-          loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(3));
-      Value bitwidth = rewriter.create<LLVM::ConstantOp>(
-          loc, rewriter.getI32Type(),
-          rewriter.getI32IntegerAttr(srcMemrefType.getElementTypeBitWidth()));
+      Value c3I32 =
+          rewriter.create<index::ConstantOp>(loc, rewriter.getIndexAttr(3));
+      Value bitwidth = rewriter.create<index::ConstantOp>(
+          loc, rewriter.getIndexAttr(srcMemrefType.getElementTypeBitWidth()));
       Value srcElementsI32 =
           rewriter.create<LLVM::TruncOp>(loc, rewriter.getI32Type(), srcBytes);
       srcBytes = rewriter.create<LLVM::LShrOp>(
           loc, rewriter.create<LLVM::MulOp>(loc, bitwidth, srcElementsI32),
           c3I32);
     }
     // Cache global (.cg) for 16 dst bytes, Cache all (.ca) for sizes other than
     // 16 dst bytes.
     NVVM::LoadCacheModifierKind cacheModifier =
         (op.getBypassL1().value_or(false) && sizeInBytes == 16)
             ? NVVM::LoadCacheModifierKind::CG
             : NVVM::LoadCacheModifierKind::CA;
 
     rewriter.create<NVVM::CpAsyncOp>(
         loc, dstPtr, scrPtr, rewriter.getI32IntegerAttr(sizeInBytes),
         NVVM::LoadCacheModifierKindAttr::get(op->getContext(), cacheModifier),
         srcBytes);
 
     // Drop the result token.
-    Value zero = rewriter.create<LLVM::ConstantOp>(
-        op->getLoc(), IntegerType::get(op.getContext(), 32),
-        rewriter.getI32IntegerAttr(0));
+    Value zero = rewriter.create<index::ConstantOp>(op->getLoc(),
+                                                    rewriter.getIndexAttr(0));
     rewriter.replaceOp(op, zero);
     return success();
   }
 };
 
 struct NVGPUAsyncCreateGroupLowering
     : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncCreateGroupOp> {
   using ConvertOpToLLVMPattern<
       nvgpu::DeviceAsyncCreateGroupOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(nvgpu::DeviceAsyncCreateGroupOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.create<NVVM::CpAsyncCommitGroupOp>(op.getLoc());
     // Drop the result token.
-    Value zero = rewriter.create<LLVM::ConstantOp>(
-        op->getLoc(), IntegerType::get(op.getContext(), 32),
-        rewriter.getI32IntegerAttr(0));
+    Value zero = rewriter.create<index::ConstantOp>(op->getLoc(),
+                                                    rewriter.getIndexAttr(0));
     rewriter.replaceOp(op, zero);
     return success();
   }
 };
 
 struct NVGPUAsyncWaitLowering
     : public ConvertOpToLLVMPattern<nvgpu::DeviceAsyncWaitOp> {
   using ConvertOpToLLVMPattern<
struct NVGPUMBarrierInitLowering
   using ConvertOpToLLVMPattern<nvgpu::MBarrierInitOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(nvgpu::MBarrierInitOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.setInsertionPoint(op);
     Value barrier = getMbarrierPtr(rewriter, *getTypeConverter(),
                                    op.getBarrier(), adaptor.getBarrier());
-
-    Value count = truncToI32(rewriter, op->getLoc(), adaptor.getCount());
+    Value count = truncToI32(rewriter, op->getLoc(), op.getCount());
 
     if (isMbarrierShared(op.getBarrier().getType())) {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierInitSharedOp>(op, barrier,
                                                               count);
     } else {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierInitOp>(op, barrier, count);
     }
     return success();
struct NVGPUMBarrierArriveNoCompleteLowering
 
   LogicalResult
   matchAndRewrite(nvgpu::MBarrierArriveNoCompleteOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value barrier = getMbarrierPtr(rewriter, *getTypeConverter(),
                                    op.getBarrier(), adaptor.getBarrier());
     Type tokenType = getTypeConverter()->convertType(
         nvgpu::MBarrierTokenType::get(op->getContext()));
-    Value count = truncToI32(rewriter, op->getLoc(), adaptor.getCount());
+    Value count = truncToI32(rewriter, op->getLoc(), op.getCount());
     if (isMbarrierShared(op.getBarrier().getType())) {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveNocompleteSharedOp>(
           op, tokenType, barrier, count);
     } else {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveNocompleteOp>(
           op, tokenType, barrier, count);
     }
     return success();
struct NVGPUMBarrierArriveExpectTxLowering
   using ConvertOpToLLVMPattern<
       nvgpu::MBarrierArriveExpectTxOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(nvgpu::MBarrierArriveExpectTxOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value barrier = getMbarrierPtr(rewriter, *getTypeConverter(),
                                    op.getBarrier(), adaptor.getBarrier());
-    Value txcount = truncToI32(rewriter, op->getLoc(), adaptor.getTxcount());
+    Value txcount = truncToI32(rewriter, op->getLoc(), op.getTxcount());
 
     if (isMbarrierShared(op.getBarrier().getType())) {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveExpectTxSharedOp>(
           op, barrier, txcount);
       return success();
     }
 
     rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveExpectTxOp>(op, barrier,
                                                                 txcount);
     return success();
   }
 };
 
 struct NVGPUMBarrierTryWaitParityLowering
     : public ConvertOpToLLVMPattern<nvgpu::MBarrierTryWaitParityOp> {
   using ConvertOpToLLVMPattern<
       nvgpu::MBarrierTryWaitParityOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(nvgpu::MBarrierTryWaitParityOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Value barrier = getMbarrierPtr(rewriter, *getTypeConverter(),
                                    op.getBarrier(), adaptor.getBarrier());
-    Value ticks = truncToI32(rewriter, op->getLoc(), adaptor.getTicks());
-    Value phase = truncToI32(rewriter, op->getLoc(), adaptor.getPhase());
+    Value ticks = truncToI32(rewriter, op->getLoc(), op.getTicks());
+    Value phase = truncToI32(rewriter, op->getLoc(), op.getPhase());
 
     if (isMbarrierShared(op.getBarrier().getType())) {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierTryWaitParitySharedOp>(
           op, barrier, phase, ticks);
       return success();
     }
 
     rewriter.replaceOpWithNewOp<NVVM::MBarrierTryWaitParityOp>(op, barrier,
                                                                phase, ticks);
     return success();
   }
 };
 
 struct NVGPUTmaAsyncLoadOpLowering
     : public ConvertOpToLLVMPattern<nvgpu::TmaAsyncLoadOp> {
   using ConvertOpToLLVMPattern<nvgpu::TmaAsyncLoadOp>::ConvertOpToLLVMPattern;
   LogicalResult
   matchAndRewrite(nvgpu::TmaAsyncLoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto dest = rewriter.create<LLVM::ExtractValueOp>(op->getLoc(),
                                                       adaptor.getDst(), 1);
     Value barrier = getMbarrierPtr(rewriter, *getTypeConverter(),
                                    op.getBarrier(), adaptor.getBarrier());
 
-    SmallVector<Value> coords = adaptor.getCoordinates();
+    SmallVector<Value> coords = op.getCoordinates();
     for (auto [index, value] : llvm::enumerate(coords)) {
       coords[index] = truncToI32(rewriter, op->getLoc(), value);
     }
 
     rewriter.replaceOpWithNewOp<NVVM::CpAsyncBulkTensorGlobalToSharedClusterOp>(
         op, dest, adaptor.getTensorMapDescriptor(), barrier, coords);
     return success();
   }
 };
 
+/// Create an i64 LLVM constant value. This should only be used with unambiguous
+/// sink operations where we know for a fact the underlying LLVM will precisely
+/// want i64.
 static Value makeI64Const(RewriterBase &rewriter, Operation *op,
                           int32_t index) {
   return rewriter.create<LLVM::ConstantOp>(op->getLoc(),
                                            rewriter.getIntegerType(64),
                                            rewriter.getI32IntegerAttr(index));
 }
 
 /// Returns a Value that holds data type enum that is expected by CUDA driver.
patterns.add<
       NVGPUTmaAsyncLoadOpLowering,           // nvgpu.tma.async.load
       NVGPUTmaCreateDescriptorOpLowering,    // nvgpu.tma.create.descriptor
       NVGPUMBarrierArriveExpectTxLowering,   // nvgpu.mbarrier.arrive.expect_tx
       NVGPUTmaAsyncLoadOpLowering,           // nvgpu.tma.async.load
       MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
       NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
       NVGPUMmaSparseSyncLowering>(converter);
 }
+
+static IntegerType getIndexTypeForMemRef(MemRefType type) {
+  if (type.getMemorySpaceAsInt() ==
+      nvgpu::NVGPUDialect::kSharedMemoryAddressSpace)
+    return IntegerType::get(type.getContext(), 32);
+  return IntegerType::get(type.getContext(), 64);
+}
+
+namespace {
+
+struct ConvertNVGPUToNVVMPass
+    : public impl::ConvertNVGPUToNVVMPassBase<ConvertNVGPUToNVVMPass> {
+  using Base::Base;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<index::IndexDialect, memref::MemRefDialect,
+                    LLVM::LLVMDialect, NVVM::NVVMDialect>();
+  }
+
+  void runOnOperation() override {
+    LowerToLLVMOptions options(&getContext());
+    options.useOpaquePointers = useOpaquePointers;
+    options.memrefIndexTypeConverter = getIndexTypeForMemRef;
+    RewritePatternSet patterns(&getContext());
+    LLVMTypeConverter converter(&getContext(), options);
+    IRRewriter rewriter(&getContext());
+    /// device-side async tokens cannot be materialized in nvvm. We just
+    /// convert them to a dummy i32 type in order to easily drop them during
+    /// conversion.
+    converter.addConversion([&](nvgpu::DeviceAsyncTokenType type) -> Type {
+      return converter.convertType(IntegerType::get(type.getContext(), 32));
+    });
+    converter.addConversion([&](nvgpu::MBarrierTokenType type) -> Type {
+      return converter.convertType(IntegerType::get(type.getContext(), 64));
+    });
+    converter.addConversion([&](nvgpu::MBarrierType type) -> Type {
+      return converter.convertType(createMBarrierMemrefType(rewriter, type));
+    });
+    converter.addConversion([&](nvgpu::TensorMapDescriptorType type) -> Type {
+      return converter.getPointerType(type.getTensor().getElementType());
+    });
+    populateNVGPUToNVVMConversionPatterns(converter, patterns);
+    LLVMConversionTarget target(getContext());
+    target.addLegalDialect<::mlir::index::IndexDialect>();
+    target.addLegalDialect<::mlir::LLVM::LLVMDialect>();
+    target.addLegalDialect<::mlir::memref::MemRefDialect>();
+    target.addLegalDialect<::mlir::NVVM::NVVMDialect>();
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+} // namespace

mlir/test/Integration/GPU/CUDA/sm90/tmaload.mlir

-// RUN: mlir-opt %s --convert-nvgpu-to-nvvm -gpu-kernel-outlining \
-// RUN:     -convert-scf-to-cf -convert-nvvm-to-llvm \
-// RUN:     -convert-vector-to-llvm \
-// RUN:     -convert-math-to-llvm \
-// RUN:     -expand-strided-metadata \
-// RUN:     -lower-affine \
-// RUN:     -convert-index-to-llvm=index-bitwidth=32 \
-// RUN:     -convert-arith-to-llvm \
-// RUN:     -finalize-memref-to-llvm \
-// RUN:     -convert-func-to-llvm \
-// RUN:     -canonicalize \
-// RUN: | mlir-opt -pass-pipeline='builtin.module(gpu.module(strip-debuginfo,convert-gpu-to-nvvm,convert-nvgpu-to-nvvm{use-opaque-pointers=1},lower-affine,convert-scf-to-cf,convert-vector-to-llvm,convert-math-to-llvm,expand-strided-metadata,lower-affine,convert-index-to-llvm{index-bitwidth=32},convert-arith-to-llvm,reconcile-unrealized-casts,gpu-to-cubin{chip=sm_90 features=+ptx80 dump-ptx}))' \
+// RUN: mlir-opt %s 
+// RUN:    -test-lower-to-nvvm="kernel-index-bitwidth=32 cubin-chip=sm_90 cubin-features=+ptx80 dump-ptx"
 // RUN: 2&>1 | FileCheck %s --check-prefixes=CHECK-PTX
 
 // CHECK-PTX: mbarrier.init.shared.b64
 // CHECK-PTX: mbarrier.arrive.expect_tx.shared.b64
 // CHECK-PTX: cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
 // CHECK-PTX: cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes
 // CHECK-PTX: mbarrier.arrive.expect_tx.shared.b64
 // CHECK-PTX: mbarrier.try_wait.parity.shared.b64
gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) threads(%arg3, %arg4, %arg5) in (%arg9 = %c128, %arg10 = %c1, %arg11 = %c1) {
         %12 = memref.load %8[%c7, %c0] : memref<8x128xf32, 3>
         gpu.printf "[GPU] TMA LOADED lhs[45][7] %f\0A" %11 : f32
         gpu.printf "[GPU] TMA LOADED rhs[7][0] %f\0A" %12 : f32
       }
       gpu.terminator
     }
     return
   }
-}
+}
 

mlir/test/lib/Dialect/GPU/TestLowerToNVVM.cpp

 #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"
 #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
 #include "mlir/Conversion/NVGPUToNVVM/NVGPUToNVVM.h"
+#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
 #include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"
 #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.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/MemRef/Transforms/Passes.h"
PassOptions::Option<std::string> cubinTriple{
       llvm::cl::init("nvptx64-nvidia-cuda")};
   PassOptions::Option<std::string> cubinChip{
       *this, "cubin-chip", llvm::cl::desc("Chip to use to serialize to cubin."),
       llvm::cl::init("sm_80")};
   PassOptions::Option<std::string> cubinFeatures{
       *this, "cubin-features",
       llvm::cl::desc("Features to use to serialize to cubin."),
       llvm::cl::init("+ptx76")};
+  PassOptions::Option<bool> dumpPtx{
+      *this, "dump-ptx", llvm::cl::desc("Whether to dump the produced ptx)"),
+      llvm::cl::init(false)};
 };
 
 //===----------------------------------------------------------------------===//
 // GPUModule-specific stuff.
 //===----------------------------------------------------------------------===//
 void buildGpuPassPipeline(OpPassManager &pm,
                           const TestLowerToNVVMOptions &options) {
   pm.addNestedPass<gpu::GPUModuleOp>(createStripDebugInfoPass());
void buildGpuPassPipeline(OpPassManager &pm,
   // Must be 64b on the host, things don't compose properly around
   // gpu::LaunchOp and gpu::HostRegisterOp.
   // TODO: fix GPU layering.
   convertIndexToLLVMPassOpt.indexBitwidth = options.kernelIndexBitWidth;
   pm.addNestedPass<gpu::GPUModuleOp>(
       createConvertIndexToLLVMPass(convertIndexToLLVMPassOpt));
 
   // TODO: C++20 designated initializers.
+  ConvertNVGPUToNVVMPassOptions convertNVGPUToNVVMPassOptions;
+  convertNVGPUToNVVMPassOptions.useOpaquePointers = true;
+  pm.addNestedPass<gpu::GPUModuleOp>(
+      createConvertNVGPUToNVVMPass(convertNVGPUToNVVMPassOptions));
+
+  pm.addNestedPass<gpu::GPUModuleOp>(createConvertSCFToCFPass());
+
+  // TODO: C++20 designated initializers.
   // The following pass is inconsistent.
   // ConvertGpuOpsToNVVMOpsOptions convertGpuOpsToNVVMOpsOptions;
   // convertGpuOpsToNVVMOpsOptions.indexBitwidth =
   //   options.kernelIndexBitWidth;
   pm.addNestedPass<gpu::GPUModuleOp>(
       // TODO: fix inconsistence.
       createLowerGpuOpsToNVVMOpsPass(/*indexBitWidth=*/
                                      options.kernelIndexBitWidth));
 
   // TODO: C++20 designated initializers.
-  ConvertNVGPUToNVVMPassOptions convertNVGPUToNVVMPassOptions;
-  convertNVGPUToNVVMPassOptions.useOpaquePointers = true;
-  pm.addNestedPass<gpu::GPUModuleOp>(
-      createConvertNVGPUToNVVMPass(convertNVGPUToNVVMPassOptions));
-  pm.addNestedPass<gpu::GPUModuleOp>(createConvertSCFToCFPass());
-
-  // TODO: C++20 designated initializers.
   GpuToLLVMConversionPassOptions gpuToLLVMConversionOptions;
   // Note: hostBarePtrCallConv must be false for now otherwise
   // gpu::HostRegister is ill-defined: it wants unranked memrefs but can't
   // lower the to bare ptr.
   gpuToLLVMConversionOptions.hostBarePtrCallConv =
       options.hostUseBarePtrCallConv;
   gpuToLLVMConversionOptions.kernelBarePtrCallConv =
       options.kernelUseBarePtrCallConv;
   gpuToLLVMConversionOptions.useOpaquePointers = true;
 
   // TODO: something useful here.
   // gpuToLLVMConversionOptions.gpuBinaryAnnotation = "";
   pm.addNestedPass<gpu::GPUModuleOp>(
       createGpuToLLVMConversionPass(gpuToLLVMConversionOptions));
 
   // Convert vector to LLVM (always needed).
   // TODO: C++20 designated initializers.
   ConvertVectorToLLVMPassOptions convertVectorToLLVMPassOptions;
   convertVectorToLLVMPassOptions.reassociateFPReductions = true;
   pm.addNestedPass<gpu::GPUModuleOp>(
       createConvertVectorToLLVMPass(convertVectorToLLVMPassOptions));
 
+  pm.addNestedPass<gpu::GPUModuleOp>(createConvertNVVMToLLVMPass());
+
   // Sprinkle some cleanups.
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
 
   // Finally we can reconcile unrealized casts.
   pm.addNestedPass<gpu::GPUModuleOp>(createReconcileUnrealizedCastsPass());
 
 #if MLIR_GPU_TO_CUBIN_PASS_ENABLE
   pm.addNestedPass<gpu::GPUModuleOp>(createGpuSerializeToCubinPass(
-      options.cubinTriple, options.cubinChip, options.cubinFeatures));
+      options.cubinTriple, options.cubinChip, options.cubinFeatures,
+      /*optLevel=*/2, /*dumpPtx=*/options.dumpPtx));
 #endif // MLIR_GPU_TO_CUBIN_PASS_ENABLE
 }
 
 void buildLowerToNVVMPassPipeline(OpPassManager &pm,
                                   const TestLowerToNVVMOptions &options) {
   //===----------------------------------------------------------------------===//
   // Host-specific stuff.
   //===----------------------------------------------------------------------===//
-  // Important, must be run at the top-level.
-  pm.addPass(createGpuKernelOutliningPass());
 
   // Important, all host passes must be run at the func level so that host
   // conversions can remain with 64 bit indices without polluting the GPU
   // kernel that may have 32 bit indices.
   // Must be 64b on the host, things don't compose properly around
   // gpu::LaunchOp and gpu::HostRegisterOp.
   // TODO: fix GPU layering.
   pm.addNestedPass<func::FuncOp>(createConvertVectorToSCFPass());
void buildLowerToNVVMPassPipeline(OpPassManager &pm,
   // TODO: fix GPU layering.
   convertFuncToLLVMPassOptions.indexBitwidth = options.hostIndexBitWidth;
   convertFuncToLLVMPassOptions.useBarePtrCallConv =
       options.hostUseBarePtrCallConv;
   convertFuncToLLVMPassOptions.useOpaquePointers = true;
   pm.addNestedPass<func::FuncOp>(
       createConvertFuncToLLVMPass(convertFuncToLLVMPassOptions));
 
-  // TODO: C++20 designated initializers.
-  ConvertIndexToLLVMPassOptions convertIndexToLLVMPassOpt;
-  // Must be 64b on the host, things don't compose properly around
-  // gpu::LaunchOp and gpu::HostRegisterOp.
-  // TODO: fix GPU layering.
-  convertIndexToLLVMPassOpt.indexBitwidth = options.hostIndexBitWidth;
-  pm.addNestedPass<func::FuncOp>(
-      createConvertIndexToLLVMPass(convertIndexToLLVMPassOpt));
-
-  pm.addNestedPass<func::FuncOp>(createArithToLLVMConversionPass());
-
   // Sprinkle some cleanups.
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(createCSEPass());
 
   //===----------------------------------------------------------------------===//
   // GPUModule-specific stuff.
   //===----------------------------------------------------------------------===//
+
+  // Due to gpu::LaunchOp and gpu::LaunchFuncOp layering and conversions, there
+  // is currently a need to call convertNVGPUToNVVM at the top-level to get
+  // proper types at the function boundary for the TMADescriptors.
+  // TODO: Fix this broken layering: conversion of TMA descriptor should be
+  // separated from introducing LLVM types.
+  // TODO: C++20 designated initializers.
+  ConvertNVGPUToNVVMPassOptions convertNVGPUToNVVMPassOptions;
+  convertNVGPUToNVVMPassOptions.useOpaquePointers = true;
+  pm.addPass(createConvertNVGPUToNVVMPass(convertNVGPUToNVVMPassOptions));
+
+  // Important, must be run at the top-level.
+  pm.addPass(createGpuKernelOutliningPass());
+
   buildGpuPassPipeline(pm, options);
 
   //===----------------------------------------------------------------------===//
   // Host post-GPUModule-specific stuff.
   //===----------------------------------------------------------------------===//
   // Convert vector to LLVM (always needed).
   // TODO: C++20 designated initializers.
   ConvertVectorToLLVMPassOptions convertVectorToLLVMPassOptions;
   convertVectorToLLVMPassOptions.reassociateFPReductions = true;
   pm.addNestedPass<func::FuncOp>(
       createConvertVectorToLLVMPass(convertVectorToLLVMPassOptions));
 
+  pm.addPass(createConvertNVVMToLLVMPass());
+
   ConvertIndexToLLVMPassOptions convertIndexToLLVMPassOpt3;
   // Must be 64b on the host, things don't compose properly around
   // gpu::LaunchOp and gpu::HostRegisterOp.
   // TODO: fix GPU layering.
   convertIndexToLLVMPassOpt3.indexBitwidth = options.hostIndexBitWidth;
   pm.addPass(createConvertIndexToLLVMPass(convertIndexToLLVMPassOpt3));
 
+  pm.addNestedPass<func::FuncOp>(createArithToLLVMConversionPass());
+
   // This must happen after cubin translation otherwise gpu.launch_func is
   // illegal if no cubin annotation is present.
   // TODO: C++20 designated initializers.
   GpuToLLVMConversionPassOptions gpuToLLVMConversionOptions;
   // Note: hostBarePtrCallConv must be false for now otherwise
   // gpu::HostRegister is ill-defined: it wants unranked memrefs but can't
   // lower the to bare ptr.
   gpuToLLVMConversionOptions.hostBarePtrCallConv =