[Mlir-commits] [mlir] [MLIR][NVVM] Add im2col_w mode support in S2G and Reduce ops (PR #214932)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 7 23:01:56 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-llvm
Author: Rajat Bajpai (rajatbajpai)
<details>
<summary>Changes</summary>
This change adds support of im2col_w mode in S2G and Reduce NVVM Dialect Ops.
---
Patch is 27.77 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/214932.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td (+3-1)
- (modified) mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp (+23-14)
- (added) mlir/test/Target/LLVMIR/nvvm/tma_store_im2colw.mlir (+37)
- (added) mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_im2col_w.mlir (+121)
- (modified) mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_invalid.mlir (+8)
``````````diff
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index ef5d276720be6..726b772f6ed68 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -4213,9 +4213,11 @@ def TMALoadModeAttr : EnumAttr<NVVM_Dialect, TMALoadMode, "tma_load_mode"> {
def TMAStoreModeTile : I32EnumAttrCase<"TILE", 0, "tile">;
def TMAStoreModeIm2Col : I32EnumAttrCase<"IM2COL", 1, "im2col">;
def TMAStoreModeTileScatter4 : I32EnumAttrCase<"TILE_SCATTER4", 2, "tile_scatter4">;
+def TMAStoreModeIm2ColW : I32EnumAttrCase<"IM2COL_W", 3, "im2col_w">;
def TMAStoreMode : I32EnumAttr<"TMAStoreMode", "NVVM TMA Store Mode",
- [TMAStoreModeTile, TMAStoreModeIm2Col, TMAStoreModeTileScatter4]> {
+ [TMAStoreModeTile, TMAStoreModeIm2Col,
+ TMAStoreModeTileScatter4, TMAStoreModeIm2ColW]> {
let genSpecializedAttr = 0;
let cppNamespace = "::mlir::NVVM";
}
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index ab51f1e5fe797..04e21e396cb43 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -137,6 +137,7 @@ LogicalResult CpAsyncBulkTensorSharedCTAToGlobalOp::verify() {
case TMAStoreMode::TILE:
return cpAsyncBulkTensorCommonVerifier(dims, false, 0, getLoc());
case TMAStoreMode::IM2COL:
+ case TMAStoreMode::IM2COL_W:
return cpAsyncBulkTensorCommonVerifier(dims, true, 0, getLoc());
case TMAStoreMode::TILE_SCATTER4:
if (dims != 5)
@@ -236,6 +237,7 @@ LogicalResult CpAsyncBulkTensorReduceOp::verify() {
case TMAStoreMode::TILE:
return cpAsyncBulkTensorCommonVerifier(dims, false, 0, getLoc());
case TMAStoreMode::IM2COL:
+ case TMAStoreMode::IM2COL_W:
return cpAsyncBulkTensorCommonVerifier(dims, true, 0, getLoc());
case TMAStoreMode::TILE_SCATTER4:
return emitError("Scatter mode unsupported for CpAsyncBulkTensorReduceOp");
@@ -4384,24 +4386,27 @@ CpAsyncBulkTensorSharedCTAToGlobalOp::getIntrinsicIDAndArgs(
args.push_back(hasCacheHint ? mt.lookupValue(cacheHint) : i64Unused);
args.push_back(builder.getInt1(hasCacheHint));
- const unsigned NI = llvm::Intrinsic::not_intrinsic;
- static constexpr llvm::Intrinsic::ID IDTable[][6] = {
- {NI, llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_1d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_2d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_3d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_4d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_5d},
- {NI, NI, NI, llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_im2col_3d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_im2col_4d,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_im2col_5d},
- {NI, NI, NI, NI, NI,
- llvm::Intrinsic::nvvm_cp_async_bulk_tensor_s2g_tile_scatter4_2d}};
+ using namespace llvm::Intrinsic;
+ const unsigned NI = not_intrinsic;
+ static constexpr ID IDTable[][6] = {
+ {NI, nvvm_cp_async_bulk_tensor_s2g_tile_1d,
+ nvvm_cp_async_bulk_tensor_s2g_tile_2d,
+ nvvm_cp_async_bulk_tensor_s2g_tile_3d,
+ nvvm_cp_async_bulk_tensor_s2g_tile_4d,
+ nvvm_cp_async_bulk_tensor_s2g_tile_5d},
+ {NI, NI, NI, nvvm_cp_async_bulk_tensor_s2g_im2col_3d,
+ nvvm_cp_async_bulk_tensor_s2g_im2col_4d,
+ nvvm_cp_async_bulk_tensor_s2g_im2col_5d},
+ {NI, NI, NI, NI, NI, nvvm_cp_async_bulk_tensor_s2g_tile_scatter4_2d},
+ {NI, NI, NI, nvvm_cp_async_bulk_tensor_s2g_im2col_w_3d,
+ nvvm_cp_async_bulk_tensor_s2g_im2col_w_4d,
+ nvvm_cp_async_bulk_tensor_s2g_im2col_w_5d}};
static_assert(getMaxEnumValForTMAStoreMode() == std::size(IDTable) - 1,
"TMAStoreModes must match number of rows in IDTable");
size_t mode = static_cast<size_t>(thisOp.getMode());
size_t dim = thisOp.getCoordinates().size();
- llvm::Intrinsic::ID id = IDTable[mode][dim];
+ ID id = IDTable[mode][dim];
if (id == llvm::Intrinsic::not_intrinsic)
llvm_unreachable(
"Invalid intrinsic for CpAsyncBulkTensorSharedCTAToGlobalOp.");
@@ -4436,7 +4441,11 @@ NVVM::IDArgPair CpAsyncBulkTensorReduceOp::getIntrinsicIDAndArgs(
nvvm_cp_async_bulk_tensor_reduce_tile_5d},
{NI, NI, NI, nvvm_cp_async_bulk_tensor_reduce_im2col_3d,
nvvm_cp_async_bulk_tensor_reduce_im2col_4d,
- nvvm_cp_async_bulk_tensor_reduce_im2col_5d}};
+ nvvm_cp_async_bulk_tensor_reduce_im2col_5d},
+ {NI, NI, NI, NI, NI, NI}, // scatter4 not supported for reduce
+ {NI, NI, NI, nvvm_cp_async_bulk_tensor_reduce_im2col_w_3d,
+ nvvm_cp_async_bulk_tensor_reduce_im2col_w_4d,
+ nvvm_cp_async_bulk_tensor_reduce_im2col_w_5d}};
size_t mode = static_cast<size_t>(thisOp.getMode());
size_t dim = thisOp.getCoordinates().size();
diff --git a/mlir/test/Target/LLVMIR/nvvm/tma_store_im2colw.mlir b/mlir/test/Target/LLVMIR/nvvm/tma_store_im2colw.mlir
new file mode 100644
index 0000000000000..9ef91604c20e0
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tma_store_im2colw.mlir
@@ -0,0 +1,37 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @tma_store_3d(%tma_desc: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_3d(ptr %0, ptr addrspace(3) %1, i32 %2, i32 %3, i32 %4, i64 %5) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.3d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i64 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.3d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i64 %5, i1 true)
+ // CHECK-NEXT: ret void
+ // CHECK-NEXT: }
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2] {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2] l2_cache_hint=%ch {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+ llvm.return
+}
+
+llvm.func @tma_store_4d(%tma_desc: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_4d(ptr %0, ptr addrspace(3) %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.4d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.4d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, i1 true)
+ // CHECK-NEXT: ret void
+ // CHECK-NEXT: }
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2,%crd3] {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2,%crd3] l2_cache_hint=%ch {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+ llvm.return
+}
+
+llvm.func @tma_store_5d(%tma_desc: !llvm.ptr, %src : !llvm.ptr<3>, %crd0: i32, %crd1: i32, %crd2: i32, %crd3: i32, %crd4: i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_5d(ptr %0, ptr addrspace(3) %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.5d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.s2g.im2col.w.5d(ptr addrspace(3) %1, ptr %0, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, i1 true)
+ // CHECK-NEXT: ret void
+ // CHECK-NEXT: }
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2,%crd3,%crd4] {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+
+ nvvm.cp.async.bulk.tensor.global.shared.cta %tma_desc, %src, box[%crd0,%crd1,%crd2,%crd3,%crd4] l2_cache_hint=%ch {mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+ llvm.return
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_im2col_w.mlir b/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_im2col_w.mlir
new file mode 100644
index 0000000000000..8d7a046c833df
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_im2col_w.mlir
@@ -0,0 +1,121 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+llvm.func @tma_store_reduce_3d_im2colw(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0 : i32, %d1 : i32, %d2 : i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_reduce_3d_im2colw(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=add */ i32 0, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=min */ i32 1, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=max */ i32 2, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=inc */ i32 3, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=dec */ i32 4, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=and */ i32 5, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=or */ i32 6, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 %5, /* red_op=xor */ i32 7, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=add */ i32 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=min */ i32 1, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=max */ i32 2, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=inc */ i32 3, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=dec */ i32 4, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=and */ i32 5, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=or */ i32 6, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.3d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i64 0, /* red_op=xor */ i32 7, i1 false)
+ // CHECK-NEXT: ret void
+ // CHECK-NEXT: }
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<add>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<min>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<max>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<inc>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<dec>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<and>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<or>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<xor>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<add>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<min>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<max>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<inc>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<dec>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<and>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<or>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2] {redKind = #nvvm.tma_redux_kind<xor>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ llvm.return
+}
+
+llvm.func @tma_store_reduce_4d_im2colw(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0 : i32, %d1 : i32, %d2 : i32, %d3 : i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_reduce_4d_im2colw(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=add */ i32 0, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=min */ i32 1, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=max */ i32 2, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=inc */ i32 3, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=dec */ i32 4, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=and */ i32 5, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=or */ i32 6, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 %6, /* red_op=xor */ i32 7, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=add */ i32 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=min */ i32 1, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=max */ i32 2, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=inc */ i32 3, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=dec */ i32 4, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=and */ i32 5, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=or */ i32 6, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.4d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i64 0, /* red_op=xor */ i32 7, i1 false)
+ // CHECK-NEXT: ret void
+ // CHECK-NEXT: }
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<add>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<min>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<max>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<inc>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<dec>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<and>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<or>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] l2_cache_hint = %ch {redKind = #nvvm.tma_redux_kind<xor>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<add>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<min>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<max>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<inc>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<dec>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<and>, mode = #nvvm.tma_store_mode<im2col_w>} : !llvm.ptr, !llvm.ptr<3>
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1, %d2, %d3] {redKind = #nvvm.tma_redux_kind<o...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/214932
More information about the Mlir-commits
mailing list