[mlir][nvgpu] Improve verifier of ldmatrix

[grypp]

PR improves the verifier of nvgpu.ldmatrix Op, so nvgpu-to-nvvm lowering does not crash.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-nvgpu

Author: Guray Ozen (grypp)

Changes

PR improves the verifier of nvgpu.ldmatrix Op, so nvgpu-to-nvvm lowering does not crash.

---

Full diff: https://github.com/llvm/llvm-project/pull/77807.diff

2 Files Affected:

• (modified) mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp (+3)
• (modified) mlir/test/Dialect/NVGPU/invalid.mlir (+8)

diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index c9756ae8fc11ce..b0a4ed1cc2697c 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -321,6 +321,9 @@ LogicalResult LdMatrixOp::verify() {
   if (isTranspose && !(elementBitWidth == 16))
     return emitError()
            << "nvgpu.ldmatrix transpose works only at 16b granularity";
+  if (resShape.size() != 2) {
+    return emitError() << "results must be 2 dimensional vector";
+  }
   if (!(resShape[1] == numElementsPer32b))
     return emitError() << "expected vector register shape[1] = "
                        << numElementsPer32b;
diff --git a/mlir/test/Dialect/NVGPU/invalid.mlir b/mlir/test/Dialect/NVGPU/invalid.mlir
index 3bffbc78569793..e1949fcfad7ad6 100644
--- a/mlir/test/Dialect/NVGPU/invalid.mlir
+++ b/mlir/test/Dialect/NVGPU/invalid.mlir
@@ -40,6 +40,14 @@ func.func @ldmatrix_trans_f32_x4(%arg0: memref<128x128xf32, 3>) ->  vector<4x1xf
 }
 // -----
 
+func.func @ldmatrix_trans_f32_x4(%arg0: memref<128x128xf32, 3>) ->  vector<4x1xf32> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{results must be 2 dimensional vector}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf32, 3> -> vector<4xf32>
+  return %a : vector<4xf32>
+}
+// -----
+
 func.func @ldmatrix_type_x4(%arg0: memref<128x128xf32, 3>) ->  vector<4x2xf16> {
   %c0  = arith.constant 0 : index
   // expected-error @+1 {{'nvgpu.ldmatrix' op failed to verify that srcMemref and res have same element type}}