[Mlir-commits] [mlir] 6ed45a0 - [MLIR][NVVM] Add tensormap.replace NVVM Op (#174926)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jan 20 21:06:46 PST 2026
Author: Srinivasa Ravi
Date: 2026-01-21T10:36:41+05:30
New Revision: 6ed45a07e86efd42fa6cb024265e5fcf659bc4d9
URL: https://github.com/llvm/llvm-project/commit/6ed45a07e86efd42fa6cb024265e5fcf659bc4d9
DIFF: https://github.com/llvm/llvm-project/commit/6ed45a07e86efd42fa6cb024265e5fcf659bc4d9.diff
LOG: [MLIR][NVVM] Add tensormap.replace NVVM Op (#174926)
This change adds the `nvvm.tensormap.replace` Op to the NVVM
dialect for the `tensormap.replace` PTX instruction.
PTX ISA Reference:
https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-tensormap-replace
Added:
mlir/test/Target/LLVMIR/nvvm/tensormap_replace.mlir
mlir/test/Target/LLVMIR/nvvm/tensormap_replace_invalid.mlir
Modified:
mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 7a45604dcc7e1..64a52acbb2278 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -6155,6 +6155,195 @@ def NVVM_Tcgen05MMAWsSparseOp : NVVM_Op<"tcgen05.mma.ws.sp",
}];
}
+//===----------------------------------------------------------------------===//
+// NVVM tensormap.replace Op
+//===----------------------------------------------------------------------===//
+
+// Tensormap Field Kind Enum
+def TensormapGlobalAddressField : I32EnumCase<"GLOBAL_ADDRESS", 0, "global_address">;
+def TensormapRankField : I32EnumCase<"RANK", 1, "rank">;
+def TensormapBoxDimField : I32EnumCase<"BOX_DIM", 2, "box_dim">;
+def TensormapGlobalDimField : I32EnumCase<"GLOBAL_DIM", 3, "global_dim">;
+def TensormapGlobalStrideField : I32EnumCase<"GLOBAL_STRIDE", 4, "global_stride">;
+def TensormapElementStrideField : I32EnumCase<"ELEMENT_STRIDE", 5, "element_stride">;
+def TensormapElemtypeField : I32EnumCase<"ELEMTYPE", 6, "elemtype">;
+def TensormapInterleaveLayoutField : I32EnumCase<"INTERLEAVE_LAYOUT", 7, "interleave_layout">;
+def TensormapSwizzleModeField : I32EnumCase<"SWIZZLE_MODE", 8, "swizzle_mode">;
+def TensormapSwizzleAtomicityField : I32EnumCase<"SWIZZLE_ATOMICITY", 9, "swizzle_atomicity">;
+def TensormapFillModeField : I32EnumCase<"FILL_MODE", 10, "fill_mode">;
+
+def TensormapField : I32Enum<"TensormapField", "NVVM Tensormap Field Kind",
+ [TensormapGlobalAddressField, TensormapRankField, TensormapBoxDimField,
+ TensormapGlobalDimField, TensormapGlobalStrideField,
+ TensormapElementStrideField, TensormapElemtypeField,
+ TensormapInterleaveLayoutField, TensormapSwizzleModeField,
+ TensormapSwizzleAtomicityField, TensormapFillModeField]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapFieldAttr : EnumAttr<NVVM_Dialect, TensormapField, "tensormap_field"> {
+ let assemblyFormat = "$value";
+}
+
+// Tensormap Elemtype Enum
+def TensormapElemtypeU8 : I32EnumCase<"U8", 0, "u8">;
+def TensormapElemtypeU16 : I32EnumCase<"U16", 1, "u16">;
+def TensormapElemtypeU32 : I32EnumCase<"U32", 2, "u32">;
+def TensormapElemtypeS32 : I32EnumCase<"S32", 3, "s32">;
+def TensormapElemtypeU64 : I32EnumCase<"U64", 4, "u64">;
+def TensormapElemtypeS64 : I32EnumCase<"S64", 5, "s64">;
+def TensormapElemtypeF16 : I32EnumCase<"F16", 6, "f16">;
+def TensormapElemtypeF32 : I32EnumCase<"F32", 7, "f32">;
+def TensormapElemtypeF32_Ftz : I32EnumCase<"F32_FTZ", 8, "f32.ftz">;
+def TensormapElemtypeF64 : I32EnumCase<"F64", 9, "f64">;
+def TensormapElemtypeBF16 : I32EnumCase<"BF16", 10, "bf16">;
+def TensormapElemtypeTF32 : I32EnumCase<"TF32", 11, "tf32">;
+def TensormapElemtypeTF32_Ftz : I32EnumCase<"TF32_FTZ", 12, "tf32.ftz">;
+def TensormapElemtypeB4x16 : I32EnumCase<"B4x16", 13, "b4x16">;
+def TensormapElemtypeB4x16P64 : I32EnumCase<"B4x16P64", 14, "b4x16_p64">;
+def TensormapElemtypeB6x16P32 : I32EnumCase<"B6x16P32", 15, "b6x16_p32">;
+
+def TensormapElemtype : I32Enum<"TensormapElemtype", "NVVM Tensormap Elemtype",
+ [TensormapElemtypeU8, TensormapElemtypeU16,
+ TensormapElemtypeU32,TensormapElemtypeS32,
+ TensormapElemtypeU64, TensormapElemtypeS64,
+ TensormapElemtypeF16, TensormapElemtypeF32,
+ TensormapElemtypeF32_Ftz,TensormapElemtypeF64,
+ TensormapElemtypeBF16, TensormapElemtypeTF32,
+ TensormapElemtypeTF32_Ftz, TensormapElemtypeB4x16,
+ TensormapElemtypeB4x16P64, TensormapElemtypeB6x16P32]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapElemtypeAttr : EnumAttr<NVVM_Dialect, TensormapElemtype, "tensormap_elemtype"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
+// Tensormap Interleave Layout Enum
+def TensormapInterleaveLayoutNoInterleave :
+ I32EnumCase<"NO_INTERLEAVE", 0, "no_interleave">;
+def TensormapInterleaveLayoutB16 : I32EnumCase<"B16", 1, "b16">;
+def TensormapInterleaveLayoutB32 : I32EnumCase<"B32", 2, "b32">;
+
+def TensormapInterleaveLayout :
+ I32Enum<"TensormapInterleaveLayout", "NVVM Tensormap Interleave Layout",
+ [TensormapInterleaveLayoutNoInterleave,
+ TensormapInterleaveLayoutB16,
+ TensormapInterleaveLayoutB32]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapInterleaveLayoutAttr : EnumAttr<NVVM_Dialect,
+ TensormapInterleaveLayout, "tensormap_interleave_layout"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
+// Tensormap Swizzle Mode Enum
+def TensormapSwizzleModeNoSwizzling : I32EnumCase<"NO_SWIZZLING", 0, "no_swizzling">;
+def TensormapSwizzleMode32B : I32EnumCase<"B32", 1, "b32">;
+def TensormapSwizzleMode64B : I32EnumCase<"B64", 2, "b64">;
+def TensormapSwizzleMode128B : I32EnumCase<"B128", 3, "b128">;
+def TensormapSwizzleMode96B : I32EnumCase<"B96", 4, "b96">;
+
+def TensormapSwizzleMode : I32Enum<"TensormapSwizzleMode", "NVVM Tensormap Swizzle Mode",
+ [TensormapSwizzleModeNoSwizzling,
+ TensormapSwizzleMode32B,
+ TensormapSwizzleMode64B,
+ TensormapSwizzleMode128B,
+ TensormapSwizzleMode96B]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapSwizzleModeAttr : EnumAttr<NVVM_Dialect,
+ TensormapSwizzleMode, "tensormap_swizzle_mode"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
+// Tensormap Swizzle Atomicity Enum
+def TensormapSwizzleAtomicityB16 : I32EnumCase<"B16", 0, "b16">;
+def TensormapSwizzleAtomicityB32 : I32EnumCase<"B32", 1, "b32">;
+def TensormapSwizzleAtomicityB32FlipB8 : I32EnumCase<"B32_FLIP_8B", 2, "b32_flip_b8">;
+def TensormapSwizzleAtomicityB64 : I32EnumCase<"B64", 3, "b64">;
+
+def TensormapSwizzleAtomicity :
+ I32Enum<"TensormapSwizzleAtomicity", "NVVM Tensormap Swizzle Atomicity",
+ [TensormapSwizzleAtomicityB16, TensormapSwizzleAtomicityB32,
+ TensormapSwizzleAtomicityB32FlipB8,
+ TensormapSwizzleAtomicityB64]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapSwizzleAtomicityAttr : EnumAttr<NVVM_Dialect,
+ TensormapSwizzleAtomicity, "tensormap_swizzle_atomicity"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
+// Tensormap Fill Mode Enum
+def TensormapFillModeZeroFill : I32EnumCase<"ZERO", 0, "zero">;
+def TensormapFillModeOOBNaNFill : I32EnumCase<"OOB_NAN", 1, "oob_nan">;
+
+def TensormapFillMode : I32Enum<"TensormapFillMode", "NVVM Tensormap Fill Mode",
+ [TensormapFillModeZeroFill, TensormapFillModeOOBNaNFill]> {
+ let cppNamespace = "::mlir::NVVM";
+}
+
+def TensormapFillModeAttr : EnumAttr<NVVM_Dialect,
+ TensormapFillMode, "tensormap_fill_mode"> {
+ let assemblyFormat = "`<` $value `>`";
+}
+
+def TensormapFieldValueAttr :
+ AnyAttrOf<[TensormapElemtypeAttr, TensormapInterleaveLayoutAttr,
+ TensormapSwizzleModeAttr, TensormapSwizzleAtomicityAttr,
+ TensormapFillModeAttr]>;
+
+def NVVM_TensormapReplaceOp : NVVM_VoidIntrinsicOp<"tensormap.replace"> {
+ let summary = "Modifies a field of the tensor-map object";
+ let description = [{
+ The `nvvm.tensormap.replace` replaces the specified field of the tensor-map
+ object at the location specified by `addr` with a new value (specified by
+ `new_value` or `new_value_attr`).
+
+ The `field` argument specifies the field of the tensor-map object to
+ replace.
+
+ `new_value` is an `i32`/`i64` argument that specifies the new value to
+ replace the `field` with for the `global_address`, `rank`, `box_dim`,
+ `global_dim`, `global_stride`, and `element_stride` fields. It must be an
+ `i64` for the `global_address` and `global_stride` fields and `i32` for the
+ remaining fields.
+
+ For `rank`, `new_value` must be one less than the desired tensor rank as
+ this field uses zero-based numbering.
+
+ `new_value_attr` is an attribute that specifies the new value to replace
+ the `field` with for the `elemtype`, `interleave_layout`, `swizzle_mode`,
+ `swizzle_atomicity`, and `fill_mode` fields. It takes the place of
+ `new_value` for these fields. It must be a valid attribute corresponding to
+ the `field` type.
+
+ The ordinal `ord` is an immediate integer argument that specifies the
+ ordinal of the `field` across the tensor which needs to be replaced and is
+ required only for the `box_dim`, `global_dim`, `global_stride`, and
+ `element_stride` fields.
+
+ [For more information, see PTX ISA.](https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-tensormap-replace)
+ }];
+
+ let hasVerifier = 1;
+ let arguments = (ins
+ TensormapFieldAttr:$field,
+ AnyTypeOf<[LLVM_PointerGlobal, LLVM_PointerShared]>:$addr,
+ Optional<AnyTypeOf<[I64, I32]>>:$new_value,
+ OptionalAttr<ConfinedAttr<I32Attr, [IntMinValue<1>, IntMaxValue<5>]>>:$ord,
+ OptionalAttr<TensormapFieldValueAttr>:$new_value_attr
+ );
+
+ let assemblyFormat = [{
+ `field` `=` $field (`[` $ord^ `]`)? `,` `new_value` `=` ($new_value_attr^):($new_value)? `in` $addr attr-dict `:` type(operands)
+ }];
+}
+
//===----------------------------------------------------------------------===//
// NVVM target attribute.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 59f9acf140074..6ce80c7456d6a 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -3006,6 +3006,82 @@ LogicalResult NVVM::ReduxOp::verify() {
return success();
}
+LogicalResult NVVM::TensormapReplaceOp::verify() {
+ auto ord = getOrd();
+ Value newVal = getNewValue();
+ auto newValAttr = getNewValueAttr();
+ auto fieldName = stringifyEnum(getField());
+
+ if (ord && !llvm::is_contained({NVVM::TensormapField::BOX_DIM,
+ NVVM::TensormapField::GLOBAL_DIM,
+ NVVM::TensormapField::GLOBAL_STRIDE,
+ NVVM::TensormapField::ELEMENT_STRIDE},
+ getField()))
+ return emitOpError("ordinal is not supported for ")
+ << fieldName << " field";
+
+ auto invalidNewVal = [&](llvm::Twine type) -> std::string {
+ return llvm::Twine("new_value must be specified and must be an " + type +
+ " for " + llvm::Twine(fieldName) + " field")
+ .str();
+ };
+
+ auto invalidNewValAttr = [&]() -> std::string {
+ return (llvm::Twine(
+ "new_value_attr must be specified and must be a valid ") +
+ llvm::Twine(fieldName) + " attribute for " + fieldName + " field")
+ .str();
+ };
+
+ switch (getField()) {
+ case NVVM::TensormapField::GLOBAL_ADDRESS:
+ if (!(newVal && newVal.getType().isInteger(64)))
+ return emitOpError(invalidNewVal("i64"));
+ break;
+ case NVVM::TensormapField::RANK:
+ if (!(newVal && newVal.getType().isInteger(32)))
+ return emitOpError(invalidNewVal("i32"));
+ break;
+ case NVVM::TensormapField::GLOBAL_STRIDE:
+ if (!ord)
+ return emitOpError("ordinal is required for global_stride field");
+ if (!(newVal && newVal.getType().isInteger(64)))
+ return emitOpError(invalidNewVal("i64"));
+ break;
+ case NVVM::TensormapField::BOX_DIM:
+ case NVVM::TensormapField::GLOBAL_DIM:
+ case NVVM::TensormapField::ELEMENT_STRIDE:
+ if (!ord)
+ return emitOpError("ordinal is required for ")
+ << stringifyEnum(getField()) << " field";
+ if (!(newVal && newVal.getType().isInteger(32)))
+ return emitOpError(invalidNewVal("i32"));
+ break;
+ case NVVM::TensormapField::ELEMTYPE:
+ if (!(newValAttr && llvm::isa<TensormapElemtypeAttr>(*newValAttr)))
+ return emitOpError(invalidNewValAttr());
+ break;
+ case NVVM::TensormapField::INTERLEAVE_LAYOUT:
+ if (!(newValAttr && llvm::isa<TensormapInterleaveLayoutAttr>(*newValAttr)))
+ return emitOpError(invalidNewValAttr());
+ break;
+ case NVVM::TensormapField::SWIZZLE_MODE:
+ if (!(newValAttr && llvm::isa<TensormapSwizzleModeAttr>(*newValAttr)))
+ return emitOpError(invalidNewValAttr());
+ break;
+ case NVVM::TensormapField::SWIZZLE_ATOMICITY:
+ if (!(newValAttr && llvm::isa<TensormapSwizzleAtomicityAttr>(*newValAttr)))
+ return emitOpError(invalidNewValAttr());
+ break;
+ case NVVM::TensormapField::FILL_MODE:
+ if (!(newValAttr && llvm::isa<TensormapFillModeAttr>(*newValAttr)))
+ return emitOpError(invalidNewValAttr());
+ break;
+ }
+
+ return success();
+}
+
/// Packs the given `field` into the `result`.
/// The `result` is 64-bits and each `field` can be 32-bits or narrower.
static llvm::Value *
@@ -4777,6 +4853,50 @@ PermuteOp::getIntrinsicIDAndArgs(Operation &op, LLVM::ModuleTranslation &mt,
return {IDs[modeIndex], args};
}
+mlir::NVVM::IDArgPair TensormapReplaceOp::getIntrinsicIDAndArgs(
+ Operation &op, LLVM::ModuleTranslation &mt, llvm::IRBuilderBase &builder) {
+ auto thisOp = cast<NVVM::TensormapReplaceOp>(op);
+
+ llvm::SmallVector<llvm::Value *> args;
+ args.push_back(mt.lookupValue(thisOp.getAddr()));
+ if (thisOp.getOrd())
+ args.push_back(builder.getInt32(thisOp.getOrd().value()));
+ if (thisOp.getNewValue())
+ args.push_back(mt.lookupValue(thisOp.getNewValue()));
+ if (auto attr = thisOp.getNewValueAttr()) {
+ auto val =
+ llvm::TypeSwitch<mlir::Attribute, unsigned>(*attr)
+ .Case<TensormapElemtypeAttr, TensormapInterleaveLayoutAttr,
+ TensormapSwizzleModeAttr, TensormapSwizzleAtomicityAttr,
+ TensormapFillModeAttr>([](auto attr) {
+ return static_cast<unsigned>(attr.getValue());
+ })
+ .Default([](auto attr) {
+ llvm_unreachable("Invalid attribute type");
+ return 0;
+ });
+ args.push_back(builder.getInt32(val));
+ }
+
+ static constexpr llvm::Intrinsic::ID IDs[] = {
+ llvm::Intrinsic::nvvm_tensormap_replace_global_address,
+ llvm::Intrinsic::nvvm_tensormap_replace_rank,
+ llvm::Intrinsic::nvvm_tensormap_replace_box_dim,
+ llvm::Intrinsic::nvvm_tensormap_replace_global_dim,
+ llvm::Intrinsic::nvvm_tensormap_replace_global_stride,
+ llvm::Intrinsic::nvvm_tensormap_replace_element_stride,
+ llvm::Intrinsic::nvvm_tensormap_replace_elemtype,
+ llvm::Intrinsic::nvvm_tensormap_replace_interleave_layout,
+ llvm::Intrinsic::nvvm_tensormap_replace_swizzle_mode,
+ llvm::Intrinsic::nvvm_tensormap_replace_swizzle_atomicity,
+ llvm::Intrinsic::nvvm_tensormap_replace_fill_mode,
+ };
+
+ unsigned fieldIndex = static_cast<unsigned>(thisOp.getField());
+
+ return {IDs[fieldIndex], args};
+}
+
//===----------------------------------------------------------------------===//
// NVVM tcgen05.mma functions
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Target/LLVMIR/nvvm/tensormap_replace.mlir b/mlir/test/Target/LLVMIR/nvvm/tensormap_replace.mlir
new file mode 100644
index 0000000000000..67576d68abb92
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tensormap_replace.mlir
@@ -0,0 +1,186 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// CHECK-LABEL: define void @tensormap_replace_global_address(ptr addrspace(1) %0, i64 %1) {
+llvm.func @tensormap_replace_global_address(%addr : !llvm.ptr<1>, %new_val : i64) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.global.address.p1(ptr addrspace(1) %0, i64 %1)
+ nvvm.tensormap.replace field = global_address, new_value = %new_val in %addr : !llvm.ptr<1>, i64
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_rank(ptr addrspace(1) %0, i32 %1) {
+llvm.func @tensormap_replace_rank(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.rank.p1(ptr addrspace(1) %0, i32 %1)
+ nvvm.tensormap.replace field = rank, new_value = %new_val in %addr : !llvm.ptr<1>, i32
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_box_dim(ptr addrspace(1) %0, i32 %1) {
+llvm.func @tensormap_replace_box_dim(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.box.dim.p1(ptr addrspace(1) %0, i32 1, i32 %1)
+ nvvm.tensormap.replace field = box_dim[1], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_global_dim(ptr addrspace(1) %0, i32 %1) {
+llvm.func @tensormap_replace_global_dim(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.global.dim.p1(ptr addrspace(1) %0, i32 1, i32 %1)
+ nvvm.tensormap.replace field = global_dim[1], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_global_stride(ptr addrspace(1) %0, i64 %1) {
+llvm.func @tensormap_replace_global_stride(%addr : !llvm.ptr<1>, %new_val : i64) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.global.stride.p1(ptr addrspace(1) %0, i32 1, i64 %1)
+ nvvm.tensormap.replace field = global_stride[1], new_value = %new_val in %addr : !llvm.ptr<1>, i64
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_element_stride(ptr addrspace(1) %0, i32 %1) {
+llvm.func @tensormap_replace_element_stride(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.element.stride.p1(ptr addrspace(1) %0, i32 1, i32 %1)
+ nvvm.tensormap.replace field = element_stride[1], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_elemtype(ptr addrspace(1) %0) {
+llvm.func @tensormap_replace_elemtype(%addr : !llvm.ptr<1>) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=u8 */ i32 0)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<u8> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=u16 */ i32 1)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<u16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=u32 */ i32 2)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<u32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=s32 */ i32 3)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<s32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=u64 */ i32 4)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<u64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=s64 */ i32 5)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<s64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=f16 */ i32 6)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<f16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=f32 */ i32 7)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<f32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=f32.ftz */ i32 8)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<f32.ftz> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=f64 */ i32 9)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<f64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=bf16 */ i32 10)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<bf16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=tf32 */ i32 11)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<tf32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=tf32.ftz */ i32 12)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<tf32.ftz> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=b4x16 */ i32 13)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<b4x16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=b4x16_p64 */ i32 14)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<b4x16_p64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.elemtype.p1(ptr addrspace(1) %0, /* elemtype=b6x16_p32 */ i32 15)
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_elemtype<b6x16_p32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_interleave_layout(ptr addrspace(1) %0) {
+llvm.func @tensormap_replace_interleave_layout(%addr : !llvm.ptr<1>) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.interleave.layout.p1(ptr addrspace(1) %0, /* interleave_layout=No interleave */ i32 0)
+ nvvm.tensormap.replace field = interleave_layout, new_value = #nvvm.tensormap_interleave_layout<no_interleave> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.interleave.layout.p1(ptr addrspace(1) %0, /* interleave_layout=16B interleave */ i32 1)
+ nvvm.tensormap.replace field = interleave_layout, new_value = #nvvm.tensormap_interleave_layout<b16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.interleave.layout.p1(ptr addrspace(1) %0, /* interleave_layout=32B interleave */ i32 2)
+ nvvm.tensormap.replace field = interleave_layout, new_value = #nvvm.tensormap_interleave_layout<b32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_swizzle_mode(ptr addrspace(1) %0) {
+llvm.func @tensormap_replace_swizzle_mode(%addr : !llvm.ptr<1>) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.mode.p1(ptr addrspace(1) %0, /* swizzle_mode=No swizzling */ i32 0)
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_mode<no_swizzling> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.mode.p1(ptr addrspace(1) %0, /* swizzle_mode=32B swizzling */ i32 1)
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_mode<b32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.mode.p1(ptr addrspace(1) %0, /* swizzle_mode=64B swizzling */ i32 2)
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_mode<b64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.mode.p1(ptr addrspace(1) %0, /* swizzle_mode=128B swizzling */ i32 3)
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_mode<b128> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.mode.p1(ptr addrspace(1) %0, /* swizzle_mode=96B swizzling */ i32 4)
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_mode<b96> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_swizzle_atomicity(ptr addrspace(1) %0) {
+llvm.func @tensormap_replace_swizzle_atomicity(%addr : !llvm.ptr<1>) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.atomicity.p1(ptr addrspace(1) %0, /* swizzle_atomicity=16B */ i32 0)
+ nvvm.tensormap.replace field = swizzle_atomicity, new_value = #nvvm.tensormap_swizzle_atomicity<b16> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.atomicity.p1(ptr addrspace(1) %0, /* swizzle_atomicity=32B */ i32 1)
+ nvvm.tensormap.replace field = swizzle_atomicity, new_value = #nvvm.tensormap_swizzle_atomicity<b32> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.atomicity.p1(ptr addrspace(1) %0, /* swizzle_atomicity=32B + 8B flip */ i32 2)
+ nvvm.tensormap.replace field = swizzle_atomicity, new_value = #nvvm.tensormap_swizzle_atomicity<b32_flip_b8> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.swizzle.atomicity.p1(ptr addrspace(1) %0, /* swizzle_atomicity=64B */ i32 3)
+ nvvm.tensormap.replace field = swizzle_atomicity, new_value = #nvvm.tensormap_swizzle_atomicity<b64> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
+
+// CHECK-LABEL: define void @tensormap_replace_fill_mode(ptr addrspace(1) %0) {
+llvm.func @tensormap_replace_fill_mode(%addr : !llvm.ptr<1>) {
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.fill.mode.p1(ptr addrspace(1) %0, /* fill_mode=Zero fill */ i32 0)
+ nvvm.tensormap.replace field = fill_mode, new_value = #nvvm.tensormap_fill_mode<zero> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: call void @llvm.nvvm.tensormap.replace.fill.mode.p1(ptr addrspace(1) %0, /* fill_mode=OOB-NaN fill */ i32 1)
+ nvvm.tensormap.replace field = fill_mode, new_value = #nvvm.tensormap_fill_mode<oob_nan> in %addr : !llvm.ptr<1>
+
+ // CHECK-NEXT: ret void
+ llvm.return
+ // CHECK-NEXT: }
+}
diff --git a/mlir/test/Target/LLVMIR/nvvm/tensormap_replace_invalid.mlir b/mlir/test/Target/LLVMIR/nvvm/tensormap_replace_invalid.mlir
new file mode 100644
index 0000000000000..245a5da40ab9f
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/nvvm/tensormap_replace_invalid.mlir
@@ -0,0 +1,129 @@
+// RUN: mlir-translate -verify-diagnostics -split-input-file -mlir-to-llvmir %s
+
+// -----
+
+llvm.func @tensormap_replace_missing_ordinal_1(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{ordinal is required for global_stride field}}
+ nvvm.tensormap.replace field = global_stride, new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_missing_ordinal_2(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{ordinal is required for box_dim field}}
+ nvvm.tensormap.replace field = box_dim, new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_ordinal(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{ordinal is not supported for rank field}}
+ nvvm.tensormap.replace field = rank[1], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_ordinal_out_of_range(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{attribute 'ord' failed to satisfy constraint: 32-bit signless integer attribute whose minimum value is 1 whose maximum value is 5}}
+ nvvm.tensormap.replace field = box_dim[6], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_missing_new_val(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value must be specified and must be an i32 for rank field}}
+ nvvm.tensormap.replace field = rank, new_value = in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_new_val_1(%addr : !llvm.ptr<1>, %new_val : i64) {
+ // expected-error @+1 {{new_value must be specified and must be an i32 for rank field}}
+ nvvm.tensormap.replace field = rank, new_value = %new_val in %addr : !llvm.ptr<1>, i64
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_new_val_2(%addr : !llvm.ptr<1>, %new_val : i64) {
+ // expected-error @+1 {{new_value must be specified and must be an i32 for box_dim field}}
+ nvvm.tensormap.replace field = box_dim[1], new_value = %new_val in %addr : !llvm.ptr<1>, i64
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_missing_new_val_attr(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid elemtype attribute for elemtype field}}
+ nvvm.tensormap.replace field = elemtype, new_value = in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_new_val_attr(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid elemtype attribute for elemtype field}}
+ nvvm.tensormap.replace field = elemtype, new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_global_address(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{new_value must be specified and must be an i64 for global_address field}}
+ nvvm.tensormap.replace field = global_address, new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_global_stride(%addr : !llvm.ptr<1>, %new_val : i32) {
+ // expected-error @+1 {{new_value must be specified and must be an i64 for global_stride field}}
+ nvvm.tensormap.replace field = global_stride[1], new_value = %new_val in %addr : !llvm.ptr<1>, i32
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_elemtype(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid elemtype attribute for elemtype field}}
+ nvvm.tensormap.replace field = elemtype, new_value = #nvvm.tensormap_interleave_layout<b16> in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_interleave_layout(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid interleave_layout attribute for interleave_layout field}}
+ nvvm.tensormap.replace field = interleave_layout, new_value = #nvvm.tensormap_swizzle_mode<b32> in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_swizzle_mode(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid swizzle_mode attribute for swizzle_mode field}}
+ nvvm.tensormap.replace field = swizzle_mode, new_value = #nvvm.tensormap_swizzle_atomicity<b32> in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_swizzle_atomicity(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid swizzle_atomicity attribute for swizzle_atomicity field}}
+ nvvm.tensormap.replace field = swizzle_atomicity, new_value = #nvvm.tensormap_fill_mode<zero> in %addr : !llvm.ptr<1>
+ llvm.return
+}
+
+// -----
+
+llvm.func @tensormap_replace_invalid_fill_mode(%addr : !llvm.ptr<1>) {
+ // expected-error @+1 {{new_value_attr must be specified and must be a valid fill_mode attribute for fill_mode field}}
+ nvvm.tensormap.replace field = fill_mode, new_value = #nvvm.tensormap_elemtype<s32> in %addr : !llvm.ptr<1>
+ llvm.return
+}
More information about the Mlir-commits
mailing list