[mlir][nvvm] Add `cp.async.bulk.tensor.shared.cluster.global`

[guraypp]

Summary

This work introduce cp.async.bulk.tensor.shared.cluster.global in NVVM dialect that executes load using TMA.

Depends on D155056

[guraypp]

rebase and fix argument order

[guraypp]

fix typos

[nicolasvasilache]

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

27

spurious include ? (should be transitively included already)

36

spurious include ?

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

string llvmBuilder = [{
     $res = createIntrinsicCall(
       builder, intId, operands);
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// NVVM TMA Ops
+//===----------------------------------------------------------------------===//
+
+def NVVM_CpAsyncBulkTensorGlobalToSharedClusterOp : NVVM_Op<"cp.async.bulk.tensor.shared.cluster.global", [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+  Arguments<(ins LLVM_i64ptr_any:$tmaDescriptor,
+                 Variadic<I32>:$coordinates,
+                 LLVM_i64ptr_shared:$dstMem,
+                 LLVM_i64ptr_shared:$mbar)> {
+  let assemblyFormat = "$tmaDescriptor `[`$coordinates `]` `,` $dstMem `,` $mbar attr-dict `:` type(operands)";
+  let extraClassDefinition = [{
+    std::string $cppClass::getPtx() {
+      int dim = getCoordinates().size();
+      std::string ptx = "cp.async.bulk.tensor.";
+      ptx += std::to_string(dim) + "d.";
+      ptx += "shared::cluster.global.mbarrier::complete_tx::bytes";
+      if(dim == 1) ptx += " [%0], [%1, {%3}\], [%2];";
+      if(dim == 2) ptx += " [%0], [%1, {%3, %4}\], [%2];";
+      if(dim == 3) ptx += " [%0], [%1, {%3, %4, %5}\], [%2];";
+      if(dim == 4) ptx += " [%0], [%1, {%3, %4, %5, %6}\], [%2];";
+      if(dim == 5) ptx += " [%0], [%1, {%3, %4, %5, %6, %7}\], [%2];";
+      return ptx;
+    }
+  }];
+  let hasVerifier = 1;
+}
+
 #endif // NVVMIR_OPS

mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp

 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/OperationSupport.h"
+#include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/AsmParser/Parser.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/SourceMgr.h"
 #include <optional>
+#include <string>
 
 using namespace mlir;
 using namespace NVVM;
 
 #include "mlir/Dialect/LLVMIR/NVVMOpsDialect.cpp.inc"
 #include "mlir/Dialect/LLVMIR/NVVMOpsEnums.cpp.inc"
 
 //===----------------------------------------------------------------------===//
return failure(parser.parseOperandList(ops) ||
                  parser.parseColonType(type) ||
                  parser.addTypeToList(type, result.types) ||
                  parser.resolveOperands(ops, {int32Ty, int1Ty},
                                         parser.getNameLoc(), result.operands));
 }
 
 void VoteBallotOp::print(OpAsmPrinter &p) { printNVVMIntrinsicOp(p, *this); }
 
+LogicalResult CpAsyncBulkTensorGlobalToSharedClusterOp::verify() {
+  if (getCoordinates().size() > 5)
+    return emitError("Maximum 5 coordinates and dimension is supported.");
+  return success();
+}
+
 LogicalResult CpAsyncOp::verify() {
   if (getModifier() != LoadCacheModifierKind::CG &&
       getModifier() != LoadCacheModifierKind::CA)
     return emitError("Only CG and CA cache modifiers are supported.");
   if (getSize() != 4 && getSize() != 8 && getSize() != 16)
     return emitError("expected byte size to be either 4, 8 or 16.");
   if (getModifier() == LoadCacheModifierKind::CG && getSize() != 16)
     return emitError("CG cache modifier is only support for 16 bytes copy.");

mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir

 // RUN: mlir-opt --convert-nvvm-to-llvm --split-input-file %s | FileCheck %s
 
 // CHECK-LABEL : @init_mbarrier_arrive_expect_tx
 llvm.func @init_mbarrier_arrive_expect_tx(%barrier : !llvm.ptr<3>, %txcount : i32) -> i64 {
   //CHECK : llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.shared.b64 $0, [$1], $2;", "=l,r,r" %{{.*}}, %{{.*}} : (!llvm.ptr<3>, i32) -> i64            
   %res = nvvm.mbarrier.arrive.expect_tx.shared %barrier, %txcount : !llvm.ptr<3>, i32 -> i64
   llvm.return %res : i64
 }
 
 // CHECK-LABEL : @init_mbarrier_arrive_expect_tx_generic
 llvm.func @init_mbarrier_arrive_expect_tx_generic(%barrier : !llvm.ptr, %txcount : i32)-> i64 {
   // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "mbarrier.arrive.expect_tx.b64 $0, [$1], $2;", "=l,l,r" %{{.*}}, %{{.*}} : (!llvm.ptr, i32) -> i64
   %res = nvvm.mbarrier.arrive.expect_tx %barrier, %txcount : !llvm.ptr, i32 -> i64
-  llvm.return %res : i64
-}
+  llvm.return %res : i64}
 
 // CHECK-LABEL : @init_mbarrier_try_wait.parity.shared
 llvm.func @init_mbarrier_try_wait_shared(%barrier : !llvm.ptr<3>, %token : i32) -> i32 {
   // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "{\0A\09.reg .pred P1; \0A\09mbarrier.try_wait.parity.shared.b64 P1, [$1], $2; \0A\09selp.b32 $0, 1, 0, P1; \0A\09}", "=r,r,r" %{{.*}}, %{{.*}} : (!llvm.ptr<3>, i32) -> i32
   %res = nvvm.mbarrier.try_wait.parity.shared %barrier, %token : !llvm.ptr<3>, i32 -> i32
   llvm.return %res : i32
 }
 

 // CHECK-LABEL : @async_cp_zfill
 func.func @async_cp_zfill(%dst: !llvm.ptr<3>, %src: !llvm.ptr<1>, %cpSize: i32) {
   // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.cg.shared.global [$0], [$1], $2, $3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
   nvvm.cp.async.shared.global %dst, %src, 16, cache =  cg, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i32
   // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.ca.shared.global [$0], [$1], $2, $3;\0A", "r,l,r" %{{.*}}, %{{.*}}, %{{.*}} : (!llvm.ptr<3>, !llvm.ptr<1>, i32) -> !llvm.void
   nvvm.cp.async.shared.global %dst, %src, 4, cache =  ca, %cpSize : !llvm.ptr<3>, !llvm.ptr<1>, i32
   return
 }
+
+// CHECK-LABEL : @tma_load_1d
+func.func @tma_load_1d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !llvm.ptr<3>, %crd0: i32) {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.bulk.tensor.1d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$3}], [$2];", "l,r,r,r"
+  nvvm.cp.async.bulk.tensor.shared.cluster.global %tmaDescriptor[%crd0], %dest, %barrier : !llvm.ptr, i32,  !llvm.ptr<3>, !llvm.ptr<3>
+  return
+}
+
+// CHECK-LABEL : @tma_load_2d
+func.func @tma_load_2d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32) {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$3, $4}], [$2];", "l,r,r,r,r"
+  nvvm.cp.async.bulk.tensor.shared.cluster.global %tmaDescriptor[%crd0, %crd1], %dest, %barrier : !llvm.ptr, i32,i32, !llvm.ptr<3>, !llvm.ptr<3>
+  return
+}
+
+// CHECK-LABEL : @tma_load_3d
+func.func @tma_load_3d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32) {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$3, $4, $5}], [$2];", "l,r,r,r,r,r"
+  nvvm.cp.async.bulk.tensor.shared.cluster.global %tmaDescriptor[%crd0, %crd1, %crd2], %dest, %barrier : !llvm.ptr, i32, i32, i32,!llvm.ptr<3>, !llvm.ptr<3>
+  return
+}
+
+// CHECK-LABEL : @tma_load_4d
+func.func @tma_load_4d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32) {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.bulk.tensor.4d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$3, $4, $5, $6}], [$2];", "l,r,r,r,r,r,r"
+  nvvm.cp.async.bulk.tensor.shared.cluster.global %tmaDescriptor[%crd0, %crd1, %crd2, %crd3], %dest, %barrier : !llvm.ptr, i32, i32,i32,i32, !llvm.ptr<3>, !llvm.ptr<3>
+  return
+}
+
+// CHECK-LABEL : @tma_load_5d
+func.func @tma_load_5d(%tmaDescriptor: !llvm.ptr, %dest : !llvm.ptr<3>, %barrier: !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32) {
+  // CHECK : llvm.inline_asm has_side_effects asm_dialect = att "cp.async.bulk.tensor.5d.shared::cluster.global.mbarrier::complete_tx::bytes [$0], [$1, {$3, $4, $5, $6, $7}], [$2];", "l,r,r,r,r,r,r,r"
+  nvvm.cp.async.bulk.tensor.shared.cluster.global %tmaDescriptor[%crd0, %crd1, %crd2, %crd3, %crd4], %dest, %barrier : !llvm.ptr, i32, i32,i32,i32,i32, !llvm.ptr<3>, !llvm.ptr<3>
+  return
+}