[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

109

typo: parameter

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

870

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

109

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

109

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

889

seems this should be prefixed with NVGPU to be consistent

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

337

typo: destination

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

575

nit: nl

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.

[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

870

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/NVGPU.td

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_TmaDescriptor : NVGPU_Type<"TmaDescriptor", "tma.descriptor", []> {
+  let summary = "TMA descriptor";
+  let parameters = (ins "MemRefType":$tensor);
+  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", [
let description = [{
     ```
   }];
   let arguments = (ins NVGPU_MBarrier:$barrier,
                        Index:$count);
   let results = (outs NVGPU_MBarrierToken:$token);
   let assemblyFormat = "$barrier `,` $count attr-dict `:` type($barrier) `->` type($token)";
 }
 
+
+def NVGPU_TmaAsyncLoadOp : NVGPU_Op<"tma.async.load", []> {
+  let summary = "TMA asynchronous load";
+  let description = [{
+    Load the give descriptor from global memory to shared memory asynchronously.
+  }];  
+  let arguments = (ins Arg<AnyMemRef, "", [MemWrite]>:$dst,
+                       NVGPU_MBarrier:$barrier,
+                       NVGPU_TmaDescriptor:$tmaDescriptor,
+                       Variadic<Index>:$coordinates);
+  let assemblyFormat = [{
+    $tmaDescriptor `[` $coordinates `]` `,` $barrier `to` $dst
+      attr-dict `:` type($tmaDescriptor) `,` type($barrier) `,` type($dst)
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // NVGPU

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::TmaDescriptorType 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();
matchAndRewrite(nvgpu::MBarrierTestWaitOp op, OpAdaptor adaptor,
     } else {
       rewriter.replaceOpWithNewOp<NVVM::MBarrierTestWaitOp>(
           op, retType, barrier, adaptor.getToken());
     }
     return success();
   }
 };
 
+struct TmaAsyncLoadOpLowering
+    : 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)) {
+      if (!value.getType().isInteger(32)) {
+        coords[index] = rewriter.create<LLVM::TruncOp>(
+            op->getLoc(), rewriter.getI32Type(), value);
+      }
+    }
+    rewriter.replaceOpWithNewOp<NVVM::CpAsyncBulkTensorGlobalToSharedClusterOp>(
+        op, dest, adaptor.getTmaDescriptor(), 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.try_wait_parity
+      TmaAsyncLoadOpLowering,                // nvgpu.tma.async.load
       MmaSyncOptoNVVM, MmaLdMatrixOpToNVVM, NVGPUAsyncCopyLowering,
       NVGPUAsyncCreateGroupLowering, NVGPUAsyncWaitLowering,
       NVGPUMmaSparseSyncLowering>(converter);
 }

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

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() << "Destionation 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_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() {
   %token = nvgpu.mbarrier.arrive.nocomplete %barrier, %count : !barrierType -> !tokenType
     
   // 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 
 }
+
+// -----
+
+// CHECK-LABEL: func @async_tma_load
+
+func.func @async_tma_load(%tmaDescriptor: !nvgpu.tma.descriptor<tensor = memref<128x128xf32,3>>, 
+                    %buffer: memref<128x128xf32, 3>, 
+                    %mbarrier : !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>) {
+  %crd0 = arith.constant 256 : index
+  %crd1 = arith.constant 256 : index
+  // CHECK : nvvm.cp.async.bulk.tensor.shared.cluster.global %{{.*}}, %{{.*}}, %{{.*}}, box[%{{.*}}, %{{.*}}] 
+  nvgpu.tma.async.load %tmaDescriptor[%crd0, %crd1], %mbarrier to %buffer : !nvgpu.tma.descriptor<tensor = memref<128x128xf32,3>>, !nvgpu.mbarrier.barrier<memorySpace = #gpu.address_space<workgroup>>, memref<128x128xf32, 3>
+  func.return 
+}