[mlir][nvgpu] Add nvgpu.tma.async.load and nvgpu.tma.descriptor

[guraypp]

Summary

This work adds nvgpu.tma.async.load Op that requests tma load asyncronusly using mbarrier object.

It also creates nvgpu.tma.descriptor type. The type is supposed be created by cuTensorMapEncodeTiled cuda drivers api.

[nicolasvasilache]

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

162

typo: parameter

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

873

some comment as to why you need the Tunc.

Also, can it ever be Ext ? (i.e. 16b or 24b addressing that we want to turn back into 32b ?)

[guraypp]

add enum types for descriptor

[guraypp]

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

162

qq: Can I use attribute here? I tried the one below, I could make tablegen compile.

DefaultValuedOptionalAttr<TensorMapSwizzleAttr, TensorMapSwizzleKind::none>:$swizzle,

[guraypp]

update on types

[nicolasvasilache]

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

162

If it works as expected, this makes sense to me (never tried to use with custom enums myself).

I suppose the benefit is that we need less boilerplate when using ?

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

939

seems this should be prefixed with NVGPU to be consistent

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

350

typo: destination

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

580

you want to expand this test a little, we should see the mbarrier ops appear and be consumed

588

Is the box rank supposed to follow the memref rank ?
Atm this looks inconsistent.

638

nit: nl

[guraypp]

add attributes

[guraypp]

address some comments

[guraypp]

fix the test

[guraypp]

add nl

[guraypp]

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

580

I need to land one more mbarriers PR. After that I can put realistic examples

588

Yes, it has to match at least for the time being. I fixed the test.

I use only base pointer of the memref.

[guraypp]

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

873

afaik, 64bits registers are emulated via 2x32bit registers. Here I truncate when we don't need 64 bit addressing.
We don't need EXT.

[guraypp]

fix enums

[guraypp]

remove using managed memory

[guraypp]

rebase

[guraypp]

rebase

[guraypp]

rebase

mlir/include/mlir/Dialect/NVGPU/IR/CMakeLists.txt

 add_mlir_dialect(NVGPU nvgpu)
 add_mlir_doc(NVGPU NVGPU Dialects/ -gen-dialect-doc)
+
+set(LLVM_TARGET_DEFINITIONS NVGPU.td)
+mlir_tablegen(NVGPUEnums.h.inc -gen-enum-decls)
+mlir_tablegen(NVGPUEnums.cpp.inc -gen-enum-defs)
+add_public_tablegen_target(MLIRNVGPUEnumsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS NVGPU.td)
+mlir_tablegen(NVGPUAttrDefs.h.inc -gen-attrdef-decls)
+mlir_tablegen(NVGPUAttrDefs.cpp.inc -gen-attrdef-defs)
+add_public_tablegen_target(MLIRNVGPUAttributesIncGen)
+
+set(LLVM_TARGET_DEFINITIONS NVGPU.td)
+mlir_tablegen(NVGPUAttrTypes.h.inc -gen-typedef-decls)
+mlir_tablegen(NVGPUAttrTypes.cpp.inc -gen-typedef-decls)
+add_public_tablegen_target(MLIRNVGPUTypesIncGen)

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

 //===----------------------------------------------------------------------===//
 
 #ifndef NVGPU
 #define NVGPU
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
+include "mlir/IR/EnumAttr.td"
 
 def NVGPU_Dialect : Dialect {
   let name = "nvgpu";
   let cppNamespace = "::mlir::nvgpu";
   let description = [{
     The `NVGPU` dialect provides a bridge between higher-level target-agnostic
     dialects (GPU and Vector) and the lower-level target-specific dialect
     (LLVM IR based NVVM dialect) for NVIDIA GPUs. This allow representing PTX
let extraClassDeclaration = [{
     /// Defines the MemRef memory space attribute numeric value that indicates
     /// a memref is located in shared memory. This should correspond to the
     /// value used in NVVM.
     static constexpr unsigned kSharedMemoryAddressSpace = 3;
   }];
 }
 
 //===----------------------------------------------------------------------===//
+// NVGPU Attribute Definitions
+//===----------------------------------------------------------------------===//
+
+def TensorMapSwizzleNone : I32EnumAttrCase<"SWIZZLE_NONE", 0, "none">;
+def TensorMapSwizzle32B  : I32EnumAttrCase<"SWIZZLE_32B", 1, "swizzle_32b">;
+def TensorMapSwizzle64B  : I32EnumAttrCase<"SWIZZLE_64B", 2, "swizzle_64b">;
+def TensorMapSwizzle128B : I32EnumAttrCase<"SWIZZLE_128B", 3, "swizzle_128b">;
+def TensorMapSwizzleKind : I32EnumAttr<"TensorMapSwizzleKind", 
+                                "Tensor map swizzling mode of shared memory banks",
+  [ TensorMapSwizzleNone, TensorMapSwizzle32B, TensorMapSwizzle64B, 
+    TensorMapSwizzle128B]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::nvgpu";
+}
+
+def TensorMapL2PromoNone : I32EnumAttrCase<"L2PROMO_NONE", 0, "none">;
+def TensorMapL2Promo64B  : I32EnumAttrCase<"L2PROMO_64B", 1, "l2promo_64b">;
+def TensorMapL2Promo128B : I32EnumAttrCase<"L2PROMO_128B", 2, "l2promo_128b">;
+def TensorMapL2Promo256B : I32EnumAttrCase<"L2PROMO_256B", 3, "l2promo_256b">;
+def TensorMapL2PromoKind : I32EnumAttr<"TensorMapL2PromoKind", 
+                                "Tensor map L2 promotion type",
+  [ TensorMapL2PromoNone, TensorMapL2Promo64B, TensorMapL2Promo128B, 
+    TensorMapL2Promo256B]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::nvgpu";
+}
+
+def TensorMapOOBZero : I32EnumAttrCase<"OOB_ZERO", 0, "zero">;
+def TensorMapOOBNaN  : I32EnumAttrCase<"OOB_NAN", 1, "nan">;
+def TensorMapOOBKind : I32EnumAttr<"TensorMapOOBKind", 
+                                "Tensor map out-of-bounds fill type",
+  [ TensorMapOOBZero, TensorMapOOBNaN]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::nvgpu";
+}
+
+def TensorMapInterleaveNone : I32EnumAttrCase<"INTERLEAVE_NONE", 0, "none">;
+def TensorMapInterleave16B  : I32EnumAttrCase<"INTERLEAVE_16B", 1, "interleave_16b">;
+def TensorMapInterleave32B  : I32EnumAttrCase<"INTERLEAVE_32B", 2, "interleave_32b">;
+def TensorMapInterleaveKind : I32EnumAttr<"TensorMapInterleaveKind", 
+                                "Tensor map interleave layout type",
+  [ TensorMapInterleaveNone, TensorMapInterleave16B, TensorMapInterleave32B]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::nvgpu";
+}
+
+def TensorMapSwizzleAttr : EnumAttr<NVGPU_Dialect, TensorMapSwizzleKind, "swizzle">;
+def TensorMapL2PromoAttr : EnumAttr<NVGPU_Dialect, TensorMapL2PromoKind, "l2promo">;
+def TensorMapOOBAttr : EnumAttr<NVGPU_Dialect, TensorMapOOBKind, "oob">;
+def TensorMapInterleaveAttr : EnumAttr<NVGPU_Dialect, TensorMapInterleaveKind, "interleave">;
+
+//===----------------------------------------------------------------------===//
 // NVGPU Type Definitions
 //===----------------------------------------------------------------------===//
 
 class NVGPU_Type<string name, string typeMnemonic,
         list<Trait> traits = []> : TypeDef<NVGPU_Dialect, name, traits> {
   let mnemonic = typeMnemonic;
 }
 
let description = [{
     https://docs.nvidia.com/cuda/parallel-thread-execution/#size-and-alignment-of-mbarrier-object
   }];    
   let parameters = (ins "Attribute":$memorySpace);
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
 def NVGPU_MBarrierToken : NVGPU_Type<"MBarrierToken", "mbarrier.token", []> { }
 
+// https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-map
+def NVGPU_TensorMapDescriptor : NVGPU_Type<"TensorMapDescriptor", "tensormap.descriptor", []> {
+  let summary = "TensorMap descriptor";
+  let parameters = (ins "MemRefType":$tensor,
+                        EnumParameter<TensorMapSwizzleKind>:$swizzle,
+                        EnumParameter<TensorMapL2PromoKind>:$l2promo,
+                        EnumParameter<TensorMapOOBKind>:$oob,
+                        EnumParameter<TensorMapInterleaveKind>:$interleave);
+  let description = [{
+    `nvgpu.tma.descriptor` is a type that represents a TMA descriptor. It is 
+    128-byte object either in constant space or kernel paramater.    
+  }];
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // NVGPU Op Definitions
 //===----------------------------------------------------------------------===//
 
 class NVGPU_Op<string mnemonic, list<Trait> traits = []> :
   Op<NVGPU_Dialect, mnemonic, traits> {}
 
 def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix", [
```mlir
       nvgpu.mbarrier.try_wait.parity %barrier, %phase, %ticks : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>
     ```
 
   }];
   let arguments = (ins NVGPU_MBarrier:$barrier, Index:$phase, Index:$ticks);
   let assemblyFormat = "$barrier `,` $phase `,` $ticks attr-dict `:` type($barrier)";  
 }
 
+def NVGPU_TmaAsyncLoadOp : NVGPU_Op<"tma.async.load", []> {
+  let summary = "TMA asynchronous load";
+  let description = [{
+    The Op loads a tile memory region from global memory to shared memory by 
+    Tensor Memory Access (TMA).
+    
+    `$tensorMapDescriptor` is tensor map descriptor which has information about
+    tile shape. The descriptor is created by `nvgpu.tma.create.descriptor`
+
+    The Op uses `$barrier` mbarrier based completion mechanism. 
+  }];  
+  let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
+                       NVGPU_MBarrier:$barrier,
+                       NVGPU_TensorMapDescriptor:$tensorMapDescriptor,
+                       Variadic<Index>:$coordinates);
+  let assemblyFormat = [{
+    $tensorMapDescriptor `[` $coordinates `]` `,` $barrier `to` $dst
+      attr-dict `:` type($tensorMapDescriptor) `,` type($barrier) `->` type($dst)
+  }];
+  let hasVerifier = 1;
+
+}
+
 #endif // NVGPU

mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h

 #define MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/Dialect/NVGPU/IR/NVGPUEnums.h.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/NVGPU/IR/NVGPUAttrDefs.h.inc"
+
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPUTypes.h.inc"
 
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.h.inc"
 
 #endif // MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

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();
struct NVGPUMBarrierInitLowering
 
   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 = adaptor.getCount();
-    if (!adaptor.getCount().getType().isInteger(32)) {
-      count = rewriter.create<LLVM::TruncOp>(op->getLoc(),
-                                             rewriter.getI32Type(), count);
-    }
+    Value count = truncToI32(rewriter, op->getLoc(), adaptor.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 = adaptor.getCount();
-    if (!adaptor.getCount().getType().isInteger(32)) {
-      count = rewriter.create<LLVM::TruncOp>(op->getLoc(),
-                                             rewriter.getI32Type(), count);
-    }
+    Value count = truncToI32(rewriter, op->getLoc(), adaptor.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
   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());
 
     if (isMbarrierShared(op.getBarrier().getType())) {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveExpectTxSharedOp>(
           op, barrier, txcount);
       return success();
     }
 
     rewriter.replaceOpWithNewOp<NVVM::MBarrierArriveExpectTxOp>(op, barrier,
                                                                 txcount);
     return success();
   }
 };
matchAndRewrite(nvgpu::MBarrierTryWaitParityOp op, OpAdaptor adaptor,
     }
 
     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();
+    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();
+  }
+};
 } // namespace
 
 void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
                                                  RewritePatternSet &patterns) {
   patterns.add<
       NVGPUMBarrierCreateLowering,           // nvgpu.mbarrier.create
       NVGPUMBarrierInitLowering,             // nvgpu.mbarrier.init
       NVGPUMBarrierArriveLowering,           // nvgpu.mbarrier.arrive
       NVGPUMBarrierArriveNoCompleteLowering, // nvgpu.mbarrier.arrive.no_complete
       NVGPUMBarrierTestWaitLowering,         // nvgpu.mbarrier.test_wait_parity
       NVGPUMBarrierTryWaitParityLowering,    // nvgpu.mbarrier.try_wait_parity
       NVGPUMBarrierArriveExpectTxLowering,   // nvgpu.mbarrier.arrive.expect_tx
+      NVGPUTmaAsyncLoadOpLowering,           // nvgpu.tma.async.load
       MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
       NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
       NVGPUMmaSparseSyncLowering>(converter);
 }

mlir/lib/Dialect/NVGPU/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRNVGPUDialect
   NVGPUDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/NVGPU
 
   DEPENDS
   MLIRNVGPUIncGen
+  MLIRNVGPUEnumsIncGen
+  MLIRNVGPUAttributesIncGen
+  MLIRNVGPUTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRGPUDialect
   MLIRIR
   MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

 //===- NVGPUDialect.cpp - MLIR NVGPU ops implementation -------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the NVGPU dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
+#include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::nvgpu;
 
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.cpp.inc"
+
 void nvgpu::NVGPUDialect::initialize() {
   addTypes<
 #define GET_TYPEDEF_LIST
 #include "mlir/Dialect/NVGPU/IR/NVGPUTypes.cpp.inc"
       >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "mlir/Dialect/NVGPU/IR/NVGPUAttrDefs.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"
       >();
 }
 
 bool nvgpu::NVGPUDialect::isSharedMemoryAddressSpace(Attribute memorySpace) {
   if (!memorySpace)
LogicalResult LdMatrixOp::verify() {
   if (!(resShape[0] == numTiles))
     return emitError()
            << "expected vector register shape[0] and numTiles to match";
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// NVGPU_TmaAsyncLoadOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult TmaAsyncLoadOp::verify() {
+  // Destination memref
+  auto dstMemref = llvm::cast<MemRefType>(getDst().getType());
+  if (!NVGPUDialect::hasSharedMemoryAddressSpace(dstMemref)) {
+    return emitError()
+           << "The operation stores data to shared memory, but "
+              "the destination memref does not have a memory space of "
+           << NVGPUDialect::kSharedMemoryAddressSpace;
+  }
+  if (getCoordinates().size() > 5) {
+    return emitError() << "Maximum 5 coordinates are supported.";
+  }
+  if (getCoordinates().size() != size_t(dstMemref.getRank())) {
+    return emitError() << "Destination memref rank is "
+                       << size_t(dstMemref.getRank()) << " but there are  "
+                       << getCoordinates().size()
+                       << " coordinates. They must match.";
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // TableGen'd dialect, type, and op definitions
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.cpp.inc"
+#define GET_ATTRDEF_CLASSES
+#include "mlir/Dialect/NVGPU/IR/NVGPUAttrDefs.cpp.inc"
+
+#include "mlir/Dialect/NVGPU/IR/NVGPUEnums.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"
 
 #define GET_TYPEDEF_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPUTypes.cpp.inc"

mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir

func.func @mbarrier_nocomplete() {
     
   // CHECK: %[[barPtr3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
   // CHECK: nvvm.mbarrier.test.wait.shared %[[barPtr3]], %[[token]]
   %isDone = nvgpu.mbarrier.test.wait %barrier, %token : !barrierType, !tokenType
 
   func.return 
 }
 
-
 // -----
 !barrierType = !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>
 !tokenType = !nvgpu.mbarrier.token
 
 // CHECK-LABEL: func @mbarrier_txcount
 func.func @mbarrier_txcount() {
       %num_threads = arith.constant 128 : index
 
     // CHECK: %[[barMemref:.+]] = memref.get_global @__mbarrier : memref<1xi64, 3>
     %barrier = nvgpu.mbarrier.create -> !barrierType
 
     // CHECK: %[[barStr:.+]] =  builtin.unrealized_conversion_cast %[[barMemref]] : memref<1xi64, 3> to !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)>
     // CHECK: %[[barPtr:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
     // CHECK: nvvm.mbarrier.init.shared %[[barPtr]]
     nvgpu.mbarrier.init %barrier, %num_threads : !barrierType
     
     %c0 = arith.constant 0 : index  
     %tidxreg = nvvm.read.ptx.sreg.tid.x : i32
     %tidx = arith.index_cast %tidxreg : i32 to index
     %cnd = arith.cmpi eq, %tidx, %c0 : index  
 
     scf.if %cnd {
       %txcount = arith.constant 256 : index
       // CHECK: %[[barPtr2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
       // CHECK: nvvm.mbarrier.arrive.expect_tx.shared %[[barPtr2]]
       nvgpu.mbarrier.arrive.expect_tx %barrier, %txcount : !barrierType
       scf.yield 
     } else {
       %txcount = arith.constant 0 : index
       // CHECK: %[[barPtr2:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
       // CHECK: nvvm.mbarrier.arrive.expect_tx.shared %[[barPtr2]]
       nvgpu.mbarrier.arrive.expect_tx %barrier, %txcount : !barrierType
       scf.yield 
     }
       
 
     %phase = arith.constant 0 : index
     %ticks = arith.constant 10000000 : index
     // CHECK: %[[barPtr3:.+]] = llvm.extractvalue %[[barStr]][1] : !llvm.struct<(ptr<3>, ptr<3>, i64, array<1 x i64>, array<1 x i64>)> 
     // CHECK: nvvm.mbarrier.try_wait.parity.shared %[[barPtr3]]
     nvgpu.mbarrier.try_wait.parity %barrier, %phase, %ticks : !barrierType
 
     func.return 
-}
+}
+
+// -----
+
+// CHECK-LABEL: func @async_tma_load
+!tensorMap1d = !nvgpu.tensormap.descriptor<tensor = memref<128xf32,3>,         swizzle=none,        l2promo = none,        oob = nan,  interleave = interleave_16b>
+!tensorMap2d = !nvgpu.tensormap.descriptor<tensor = memref<32x32xf32,3>,       swizzle=swizzle_32b, l2promo = none,        oob = zero, interleave = none>
+!tensorMap3d = !nvgpu.tensormap.descriptor<tensor = memref<2x32x32xf32,3>,     swizzle=swizzle_64b, l2promo = l2promo_64b, oob = zero, interleave = none>
+!tensorMap4d = !nvgpu.tensormap.descriptor<tensor = memref<2x2x32x32xf32,3>,   swizzle=swizzle_128b,l2promo = l2promo_128b,oob = zero, interleave = none>
+!tensorMap5d = !nvgpu.tensormap.descriptor<tensor = memref<2x2x2x32x32xf32,3>, swizzle=none,        l2promo = none,        oob = zero, interleave = none>
+!mbarrier = !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>
+func.func @async_tma_load(%tensorMap1d: !tensorMap1d, %tensorMap2d: !tensorMap2d, %tensorMap3d: !tensorMap3d, %tensorMap4d: !tensorMap4d, %tensorMap5d: !tensorMap5d, 
+                              %buffer1d: memref<128xf32,3>,      
+                              %buffer2d: memref<32x32xf32,3>,    
+                              %buffer3d: memref<2x32x32xf32,3>,  
+                              %buffer4d: memref<2x2x32x32xf32,3>,  
+                              %buffer5d: memref<2x2x2x32x32xf32,3>,
+                              %mbarrier: !mbarrier) {
+  %crd0 = arith.constant 0 : index
+  %crd1 = arith.constant 0 : index
+  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}] 
+  nvgpu.tma.async.load %tensorMap1d[%crd0], %mbarrier to %buffer1d : !tensorMap1d, !mbarrier -> memref<128xf32,3>
+  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}] 
+  nvgpu.tma.async.load %tensorMap2d[%crd0, %crd1], %mbarrier to %buffer2d : !tensorMap2d, !mbarrier -> memref<32x32xf32,3>
+  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}] 
+  nvgpu.tma.async.load %tensorMap3d[%crd0, %crd1, %crd0], %mbarrier to %buffer3d : !tensorMap3d, !mbarrier -> memref<2x32x32xf32,3>
+  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] 
+  nvgpu.tma.async.load %tensorMap4d[%crd0, %crd1, %crd1, %crd0], %mbarrier to %buffer4d : !tensorMap4d, !mbarrier -> memref<2x2x32x32xf32,3>
+  // CHECK: nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] 
+  nvgpu.tma.async.load %tensorMap5d[%crd0, %crd1, %crd1, %crd0, %crd0], %mbarrier to %buffer5d : !tensorMap5d, !mbarrier -> memref<2x2x2x32x32xf32,3>
+  func.return 
+}
+