[Mlir-commits] [mlir] [mlir][nvvm] Add prefetch.tensormap (PR #67564)
Guray Ozen
llvmlistbot at llvm.org
Tue Oct 17 03:58:29 PDT 2023
https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/67564
>From 9a09ec563ff2a56f6468e0e4cbac552e2e85edcc Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Wed, 27 Sep 2023 16:51:38 +0200
Subject: [PATCH 1/5] [mlir][nvvm] Add prefetch.tensormap
This PR adds `prefetch.tensormap` Op. It brings the cache line containing the given tma descriptor for subsequent use by the cp.async.bulk.tensor instruction.
https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-prefetch-prefetchu
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 15 +++++++++++++++
mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir | 11 +++++++++++
2 files changed, 26 insertions(+)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index d550fe1f33140ed..b1176d517d05600 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1438,6 +1438,21 @@ def NVVM_CpAsyncBulkTensorSharedCTAToGlobalOp :
let hasVerifier = 1;
}
+def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap",
+ [DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
+ Arguments<(ins LLVM_i64ptr_any:$tmaDescriptor, PtxPredicate:$predicate)> {
+ let description = [{
+ The Op brings the cache line containing the given $tmaDescriptor for
+ subsequent use by the `cp.async.bulk.tensor` instruction.
+ }];
+ let assemblyFormat = "$tmaDescriptor (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
+ let extraClassDefinition = [{
+ std::string $cppClass::getPtx() {
+ return std::string("prefetch.tensormap [%0];");
+ }
+ }];
+}
+
//===----------------------------------------------------------------------===//
// NVVM Wgmma Ops
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
index fcc882f562a4a95..0d0ac9637438a95 100644
--- a/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
+++ b/mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
@@ -504,3 +504,14 @@ func.func @elect_one_leader_sync() {
%cnd = nvvm.elect.sync -> i1
return
}
+
+// -----
+
+// CHECK-LABEL: @init_mbarrier_arrive_expect_tx
+llvm.func @init_mbarrier_arrive_expect_tx(%desc : !llvm.ptr, %pred : i1) {
+ //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "prefetch.tensormap [$0];", "l"
+ nvvm.prefetch.tensormap %desc : !llvm.ptr
+ //CHECK: llvm.inline_asm has_side_effects asm_dialect = att "@$1 prefetch.tensormap [$0];", "l,b"
+ nvvm.prefetch.tensormap %desc, predicate = %pred : !llvm.ptr, i1
+ llvm.return
+}
>From 8e79371ee0cfa0167bc9ba87f031308f1ffe6ccf Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Wed, 27 Sep 2023 16:58:53 +0200
Subject: [PATCH 2/5] implement it in nvgpu dialect as well
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 4 ----
mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td | 12 ++++++++++++
mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp | 13 +++++++++++++
3 files changed, 25 insertions(+), 4 deletions(-)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index b1176d517d05600..cefdd7cc4033a11 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -1441,10 +1441,6 @@ def NVVM_CpAsyncBulkTensorSharedCTAToGlobalOp :
def NVVM_PrefetchTensorMapOp : NVVM_Op<"prefetch.tensormap",
[DeclareOpInterfaceMethods<BasicPtxBuilderOpInterface>]>,
Arguments<(ins LLVM_i64ptr_any:$tmaDescriptor, PtxPredicate:$predicate)> {
- let description = [{
- The Op brings the cache line containing the given $tmaDescriptor for
- subsequent use by the `cp.async.bulk.tensor` instruction.
- }];
let assemblyFormat = "$tmaDescriptor (`,` `predicate` `=` $predicate^)? attr-dict `:` type(operands)";
let extraClassDefinition = [{
std::string $cppClass::getPtx() {
diff --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
index fd16376be366912..dd00355b6d77e33 100644
--- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
+++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
@@ -619,6 +619,18 @@ def NVGPU_MBarrierTryWaitParityOp : NVGPU_Op<"mbarrier.try_wait.parity", []> {
let assemblyFormat = "$barriers `[` $mbarId `]` `,` $phase `,` $ticks attr-dict `:` type($barriers)";
}
+def NVGPU_TmaPrefetchOp : NVGPU_Op<"tma.prefetch.descriptor", []> {
+ let summary = "Prefetch given `nvgpu.tensormap.descriptor` ";
+ let description = [{
+ The Op brings the cache line containing the given `$tmaDescriptor` for
+ subsequent use by the `tma.async.load` instruction.
+ }];
+ let arguments = (ins NVGPU_TensorMapDescriptor:$tensorMapDescriptor, Optional<I1>:$predicate);
+ let assemblyFormat = [{
+ $tensorMapDescriptor (`,` $predicate^)? attr-dict `:` type($tensorMapDescriptor)
+ }];
+}
+
def NVGPU_TmaAsyncLoadOp : NVGPU_Op<"tma.async.load", []> {
let summary = "TMA asynchronous load";
let description = [{
diff --git a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
index 029659a2f855416..7eb6f42d2788e35 100644
--- a/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
+++ b/mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
@@ -1610,6 +1610,18 @@ struct NVGPUWarpgroupMmaInitAccumulatorOpLowering
}
};
+struct NVGPUTmaPrefetchOpLowering
+ : public ConvertOpToLLVMPattern<nvgpu::TmaPrefetchOp> {
+ using ConvertOpToLLVMPattern<nvgpu::TmaPrefetchOp>::ConvertOpToLLVMPattern;
+ LogicalResult
+ matchAndRewrite(nvgpu::TmaPrefetchOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ rewriter.replaceOpWithNewOp<NVVM::PrefetchTensorMapOp>(
+ op, adaptor.getTensorMapDescriptor(), adaptor.getPredicate());
+ return success();
+ }
+};
+
} // namespace
void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
@@ -1623,6 +1635,7 @@ void mlir::populateNVGPUToNVVMConversionPatterns(LLVMTypeConverter &converter,
NVGPUMBarrierTryWaitParityLowering, // nvgpu.mbarrier.try_wait_parity
NVGPUTmaAsyncLoadOpLowering, // nvgpu.tma.async.load
NVGPUTmaCreateDescriptorOpLowering, // nvgpu.tma.create.descriptor
+ NVGPUTmaPrefetchOpLowering, // nvgpu.tma.prefetch.descriptor
NVGPUMBarrierArriveExpectTxLowering, // nvgpu.mbarrier.arrive.expect_tx
NVGPUGenerateWarpgroupDescriptorLowering, // nvgpu.warpgroup.generate.descriptor
NVGPUWarpgroupMmaOpLowering, // nvgpu.warpgroup.mma
>From 1c07e75e80d752c47fa79c2e528b30f0fa62a739 Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Wed, 27 Sep 2023 17:05:26 +0200
Subject: [PATCH 3/5] test
---
mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir | 7 +++++++
1 file changed, 7 insertions(+)
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
index bf660e2683158e5..bd9c12ca5d9b408 100644
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
@@ -643,6 +643,13 @@ func.func @create_tensor_map(%devicePtr2d : memref<64x128xf32>, %devicePtr1d : m
func.return
}
+func.func @prefetch(%tensorMap1d: !tensorMap1d, %p : i1) {
+ // CHECK: nvvm.prefetch.tensormap
+ nvgpu.tma.prefetch.descriptor %tensorMap1d: !tensorMap1d
+ // CHECK: nvvm.prefetch.tensormap
+ nvgpu.tma.prefetch.descriptor %tensorMap1d, %p: !tensorMap1d
+}
+
!lhsTensorMap = !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = none, oob = zero, interleave = none>
!rhsTensorMap = !nvgpu.tensormap.descriptor<tensor = memref<64x128xf16, strided<[128, 1], offset: 8192>, 3>, swizzle = swizzle_128b, l2promo = none, oob = zero, interleave = none>
>From 78f5a0328dddf96a3bbf77c3571709098cd4622f Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Thu, 28 Sep 2023 08:31:55 +0200
Subject: [PATCH 4/5] check the arguments in the test
---
mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
index bd9c12ca5d9b408..2b826dc153836e7 100644
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
@@ -643,10 +643,12 @@ func.func @create_tensor_map(%devicePtr2d : memref<64x128xf32>, %devicePtr1d : m
func.return
}
-func.func @prefetch(%tensorMap1d: !tensorMap1d, %p : i1) {
- // CHECK: nvvm.prefetch.tensormap
+// CHECK-LABEL: @tma_prefetch(
+// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.tensormap.descriptor<tensor = memref<128xf32,3>, swizzle=none, l2promo = none, oob = nan, interleave = none>, %[[arg1:[a-zA-Z0-9_]+]]: i1
+func.func @tma_prefetch(%tensorMap1d: !tensorMap1d, %p : i1) {
+ // CHECK: nvvm.prefetch.tensormap %[[arg0]] : !llvm.ptr
nvgpu.tma.prefetch.descriptor %tensorMap1d: !tensorMap1d
- // CHECK: nvvm.prefetch.tensormap
+ // CHECK: nvvm.prefetch.tensormap %[[arg0]], predicate = %[[arg1]] : !llvm.ptr, i1
nvgpu.tma.prefetch.descriptor %tensorMap1d, %p: !tensorMap1d
}
>From b18725e43976db61452bb0007fe4370f72286944 Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Tue, 17 Oct 2023 12:58:11 +0200
Subject: [PATCH 5/5] fix the test
---
mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir | 8 +++++---
1 file changed, 5 insertions(+), 3 deletions(-)
diff --git a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
index 2b826dc153836e7..8971585e03c7add 100644
--- a/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
+++ b/mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
@@ -644,12 +644,14 @@ func.func @create_tensor_map(%devicePtr2d : memref<64x128xf32>, %devicePtr1d : m
}
// CHECK-LABEL: @tma_prefetch(
-// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.tensormap.descriptor<tensor = memref<128xf32,3>, swizzle=none, l2promo = none, oob = nan, interleave = none>, %[[arg1:[a-zA-Z0-9_]+]]: i1
+// CHECK-SAME: %[[arg0:[a-zA-Z0-9_]+]]: !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none>, %[[arg1:[a-zA-Z0-9_]+]]: i1
func.func @tma_prefetch(%tensorMap1d: !tensorMap1d, %p : i1) {
- // CHECK: nvvm.prefetch.tensormap %[[arg0]] : !llvm.ptr
+ // CHECK: %[[S0:.+]] = builtin.unrealized_conversion_cast %[[arg0]] : !nvgpu.tensormap.descriptor<tensor = memref<128xf32, 3>, swizzle = none, l2promo = none, oob = nan, interleave = none> to !llvm.ptr
+ // CHECK: nvvm.prefetch.tensormap %[[S0]] : !llvm.ptr
nvgpu.tma.prefetch.descriptor %tensorMap1d: !tensorMap1d
- // CHECK: nvvm.prefetch.tensormap %[[arg0]], predicate = %[[arg1]] : !llvm.ptr, i1
+ // CHECK: nvvm.prefetch.tensormap %[[S0]], predicate = %[[arg1]] : !llvm.ptr, i1
nvgpu.tma.prefetch.descriptor %tensorMap1d, %p: !tensorMap1d
+ func.return
}
!lhsTensorMap = !nvgpu.tensormap.descriptor<tensor = memref<128x64xf16, 3>, swizzle = swizzle_128b, l2promo = none, oob = zero, interleave = none>
More information about the Mlir-commits
mailing list