[Mlir-commits] [mlir] [MLIR][NVVM] Add im2col_w mode support in S2G and Reduce ops (PR #214932)
Rajat Bajpai
llvmlistbot at llvm.org
Fri Aug 7 23:01:16 PDT 2026
https://github.com/rajatbajpai created https://github.com/llvm/llvm-project/pull/214932
This change adds support of im2col_w mode in S2G and Reduce NVVM Dialect Ops.
>From 0ab4342ebc75cf96d2521e5f4923f0833b476df7 Mon Sep 17 00:00:00 2001
From: Rajat Bajpai <rbajpai at nvidia.com>
Date: Fri, 7 Aug 2026 18:38:29 +0000
Subject: [PATCH] [MLIR][NVVM] Add im2col_w mode support in S2G and Reduce ops
This change adds support of im2col_w mode in S2G and Reduce NVVM Dialect
Ops.
---
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td | 4 +-
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 37 ++++--
.../Target/LLVMIR/nvvm/tma_store_im2colw.mlir | 37 ++++++
.../nvvm/tma_store_reduce_im2col_w.mlir | 121 ++++++++++++++++++
.../LLVMIR/nvvm/tma_store_reduce_invalid.mlir | 8 ++
5 files changed, 192 insertions(+), 15 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/nvvm/tma_store_im2colw.mlir
create mode 100644 mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_im2col_w.mlir
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<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] {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_5d_im2colw(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0 : i32, %d1 : i32, %d2 : i32, %d3 : i32, %d4 : i32, %ch : i64) {
+ // CHECK-LABEL: define void @tma_store_reduce_5d_im2colw(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7) {
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=add */ i32 0, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=min */ i32 1, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=max */ i32 2, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=inc */ i32 3, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=dec */ i32 4, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=and */ i32 5, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=or */ i32 6, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 %7, /* red_op=xor */ i32 7, i1 true)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=add */ i32 0, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=min */ i32 1, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=max */ i32 2, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=inc */ i32 3, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=dec */ i32 4, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=and */ i32 5, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, i64 0, /* red_op=or */ i32 6, i1 false)
+ // CHECK-NEXT: call void @llvm.nvvm.cp.async.bulk.tensor.reduce.im2col.w.5d(ptr addrspace(3) %0, ptr %1, i32 %2, i32 %3, i32 %4, i32 %5, i32 %6, 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, %d4] 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, %d4] 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, %d4] 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, %d4] 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, %d4] 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, %d4] 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, %d4] 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, %d4] 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, %d4] {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, %d4] {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, %d4] {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, %d4] {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, %d4] {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, %d4] {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, %d4] {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, %d4] {redKind = #nvvm.tma_redux_kind<xor>, 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_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_invalid.mlir
index 2fcf00fa3b670..577409916b0d6 100644
--- a/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvm/tma_store_reduce_invalid.mlir
@@ -18,6 +18,14 @@ llvm.func @tma_reduce_2d_im2col(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0
// -----
+llvm.func @tma_reduce_2d_im2col_w(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0 : i32, %d1 : i32, %ch : i64) {
+ // expected-error @below {{to use im2col mode, the tensor has to be at least 3-dimensional}}
+ nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0, %d1] {redKind = #nvvm.tma_redux_kind<and>, mode = #nvvm.tma_store_mode<im2col_w>}: !llvm.ptr, !llvm.ptr<3>
+ llvm.return
+}
+
+// -----
+
llvm.func @tma_store_reduce_scatter(%src : !llvm.ptr<3>, %tma_desc : !llvm.ptr, %d0 : i32, %ch : i64) {
// expected-error @below {{Scatter mode unsupported for CpAsyncBulkTensorReduceOp}}
nvvm.cp.async.bulk.tensor.reduce %tma_desc, %src, box[%d0] {redKind = #nvvm.tma_redux_kind<add>, mode = #nvvm.tma_store_mode<tile_scatter4>} : !llvm.ptr, !llvm.ptr<3>
More information about the Mlir-commits
mailing list