[Mlir-commits] [mlir] [MLIR][NVVM] Update Op verifiers to prevent ungraceful exits (PR #165677)

[Srinivasa Ravi]

https://github.com/Wolfram70 updated https://github.com/llvm/llvm-project/pull/165677

>From f566fa78261aab88fcb77f06437b53119e52fb3f Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Thu, 30 Oct 2025 08:27:40 +0000
Subject: [PATCH 1/2] [MLIR][NVVM] Update Op verifiers to prevent ungraceful
 exits

Updates the following Ops to prevent ungraceful exits with a
stack-dump in certain cases of incorrect usages, and instead
gracefully error out with a more informative error message:

- tcgen05.ld
- shfl.sync
---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp  | 24 ++++++++++++++-------
 mlir/test/Dialect/LLVMIR/invalid.mlir       |  7 ++++++
 mlir/test/Target/LLVMIR/nvvmir-invalid.mlir |  8 +++++++
 3 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index f0de4dbcc1d4b..402c90fba0f2d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -867,15 +867,20 @@ LogicalResult MmaOp::verify() {
 }
 
 LogicalResult ShflOp::verify() {
-  if (!(*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid"))
-    return success();
   auto type = llvm::dyn_cast<LLVM::LLVMStructType>(getType());
-  auto elementType = (type && type.getBody().size() == 2)
-                         ? llvm::dyn_cast<IntegerType>(type.getBody()[1])
-                         : nullptr;
-  if (!elementType || elementType.getWidth() != 1)
-    return emitError("expected return type to be a two-element struct with "
-                     "i1 as the second element");
+
+  if ((*this)->getAttrOfType<UnitAttr>("return_value_and_is_valid")) {
+    auto elementType = (type && type.getBody().size() == 2)
+                           ? llvm::dyn_cast<IntegerType>(type.getBody()[1])
+                           : nullptr;
+    if (!elementType || elementType.getWidth() != 1)
+      return emitOpError("expected return type to be a two-element struct with "
+                         "i1 as the second element");
+  } else {
+    if (type)
+      return emitOpError("\"return_value_and_is_valid\" attribute must be "
+                         "specified when returning the predicate");
+  }
   return success();
 }
 
@@ -2450,6 +2455,9 @@ LogicalResult Tcgen05LdOp::verify() {
   LogicalResult result = success();
   if (getShape() == NVVM::Tcgen05LdStShape::SHAPE_16X32BX2 && !getOffset())
     result = emitError("shape 16x32bx2 requires offset argument");
+  
+  if (getShape() != NVVM::Tcgen05LdStShape::SHAPE_16X32BX2 && getOffset())
+    result = emitError("offset argument is only supported for shape 16x32bx2");
 
   auto resTy = getRes().getType();
   unsigned resLen = isa<VectorType>(resTy)
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index aaf9f8024bfbe..90208aa55bd55 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -684,6 +684,13 @@ func.func @nvvm_invalid_shfl_pred_3(%arg0 : i32, %arg1 : i32, %arg2 : i32, %arg3
 
 // -----
 
+func.func @nvvm_invalid_shfl_pred_4(%arg0 : i32, %arg1 : f32, %arg2 : i32, %arg3 : i32) {
+  // expected-error at +1 {{"return_value_and_is_valid" attribute must be specified when returning the predicate}}
+  %0 = nvvm.shfl.sync bfly %arg0, %arg1, %arg2, %arg3 : f32 -> !llvm.struct<(f32, i1)>
+}
+
+// -----
+
 func.func @nvvm_invalid_mma_0(%a0 : f16, %a1 : f16,
                          %b0 : vector<2xf16>, %b1 : vector<2xf16>,
                          %c0 : f32, %c1 : f32, %c2 : f32, %c3 : f32,
diff --git a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
index 09b8f593154b5..8cb7b068498fd 100644
--- a/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir-invalid.mlir
@@ -621,3 +621,11 @@ func.func @invalid_range_equal_bounds() {
   %0 = nvvm.read.ptx.sreg.warpsize range <i32, 32, 32> : i32
   return
 }
+
+// -----
+
+llvm.func @nvvm_tcgen05_ld_32x32b_offset(%tmemAddr : !llvm.ptr<6>, %offset : i64) -> () {
+  // expected-error at +1 {{offset argument is only supported for shape 16x32bx2}}
+  %ldv2 = nvvm.tcgen05.ld %tmemAddr, %offset { pack, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<2 x i32>
+  llvm.return
+}

>From 045ce6c99786cb8634ca885dfca5b94a58ce2b8a Mon Sep 17 00:00:00 2001
From: Srinivasa Ravi <srinivasar at nvidia.com>
Date: Thu, 30 Oct 2025 08:37:55 +0000
Subject: [PATCH 2/2] fix formatting

---
 mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index 402c90fba0f2d..a23245f92cee7 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -2455,7 +2455,7 @@ LogicalResult Tcgen05LdOp::verify() {
   LogicalResult result = success();
   if (getShape() == NVVM::Tcgen05LdStShape::SHAPE_16X32BX2 && !getOffset())
     result = emitError("shape 16x32bx2 requires offset argument");
-  
+
   if (getShape() != NVVM::Tcgen05LdStShape::SHAPE_16X32BX2 && getOffset())
     result = emitError("offset argument is only supported for shape 16x32bx2");