[mlir][NVGPU] Verifier for nvgpu.ldmatrix

[manishucsd]

Summary

• Adds verifiers for nvgpu.ldmatrix op
• Adds tests to mlir/test/Dialect/NVGPU/invalid.mlir

[ThomasRaoux]

Looks good, couple nit comments

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

55–60

nit: can you improve the formatting? (clang-format doesn't work on tablegen files :( )

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

228

nit: skip else after return

228

forgot to add the coding standard associated link:
https://llvm.org/docs/CodingStandards.html#don-t-use-else-after-a-return

[manishucsd]

• Addressing review comments

[ThomasRaoux]

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

198

one more nit: comments should be a full sentence meaning they should end with a period.
https://llvm.org/docs/CodingStandards.html#commenting

[manishucsd]

cleaning comments

[manishucsd]

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

198

addressed llvm-style commenting guideline while maintaining comment-style consistency with the rest of the file

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

 
 
 //===----------------------------------------------------------------------===//
 // NVGPU Op definitions
 //===----------------------------------------------------------------------===//
 
 class NVGPU_Op<string mnemonic, list<Trait> traits = []> :
   Op<NVGPU_Dialect, mnemonic, traits> {}
 
-def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix",
-                                [MemoryEffects<[MemRead]>]> {
+def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix", [
+                                MemoryEffects<[MemRead]>,
+                                PredOpTrait<"srcMemref and res have same element type", 
+                                            TCresVTEtIsSameAsOp<0, 0>>]> {
   let description = [{
   The `nvgpu.ldmatrix` op represents loading a matrix fragment from
   memory. The load source and result type must be compatible with lowering
   to the `nvvm.ldmatrix` instruction. This op is meant to represent
   the distributed version of a `vector.transfer_read` as an intermediate
   step between lowering from `vector.transfer_read` to `nvvm.ldmatrix`.
 
   This operation is meant to follow the semantic of described here:
   https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#warp-level-matrix-instructions-ldmatrix
 
   Example:
   ```mlir
   %0 = nvgpu.ldmatrix %sm[%c0, %c0] {numTiles = 4 : i32, transpose = false} :
     memref<?x?xf16, 3> -> vector<4x2xf16>
   ```
   }];
 
   let arguments = (ins Arg<AnyMemRef, "", [MemRead]>:$srcMemref,
                            Variadic<Index>:$indices, BoolAttr:$transpose,
                            I32Attr:$numTiles);
   let results = (outs AnyVector:$res);
   let assemblyFormat = [{
     $srcMemref`[` $indices `]` attr-dict `:` type($srcMemref) `->` type($res)
   }];
+
+  let hasVerifier = 1;
 }
 
 def NVGPU_MmaSyncOp : NVGPU_Op<"mma.sync", [
                                NoSideEffect,
-    PredOpTrait<"matrixA and matrixB have same element type", TCopVTEtIsSameAs<0, 1>>,
-  ]> {
+                               PredOpTrait<"matrixA and matrixB have same element type", 
+                                           TCopVTEtIsSameAs<0, 1>>]> {
   let description = [{
   The `nvgpu.mma.sync` op represents the distributed form of a collective
   matrix-multiply-and-accumulate (mma) operation that is compatible with
   `nvvm.mma.sync`. The operands and results are fragments of the full matrix
   operands. The full shape of the distributed mma operation is given by the
   `mmaShape` attribute in the form of a list of dimensions `[m, n, k]`.
 
   This operation is meant to be lowered to the `nvvm.mma.sync` instruction, and
let assemblyFormat = [{
     `(` $matrixA`,` $matrixB`,` $matrixC `)` attr-dict
     `:` `(` type($matrixA) `,` type($matrixB) `,` type($matrixC) `)` `->` type($res)
   }];
 
   let hasVerifier = 1;
 }
 
 
-def NVGPU_DeviceAsyncCopyOp : NVGPU_Op<"device_async_copy",
-  [AttrSizedOperandSegments]> {
+def NVGPU_DeviceAsyncCopyOp : NVGPU_Op<"device_async_copy", [
+                                       AttrSizedOperandSegments]> {
   let summary = "device-side asynchronous copy";
   let description = [{
     The `gpu.device_async_copy` op initiates an asynchronous copy operation of
     `$size` elements from source to the destination without blocking the thread.
     The destination has to be in shared memory.
 
     This is memory access will be pending to be added to a group.
 

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

return emitOpError() << "expected " << srcMemref.getRank()
                          << " source indices, got " << getSrcIndices().size();
   if (size_t(dstMemref.getRank()) != getDstIndices().size())
     return emitOpError() << "expected " << dstMemref.getRank()
                          << " destination indices, got "
                          << getDstIndices().size();
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// NVGPU_MmaSyncOp
+//===----------------------------------------------------------------------===//
+
 LogicalResult MmaSyncOp::verify() {
 
   // Fundamental tensor core mma.sync op
   // For F32 (TF32), F16, S8, and S4 data types fundamental tensor core
   // operation is of shape: 8-by-8-by-128b. F64 is an exception. The
   // verification for mma.sync covering various shapes and data types is based
   // on the fundamental tensor core operionation.
   constexpr int kThreads = 32; // 32 threads per warp
LogicalResult MmaSyncOp::verify() {
   // verify shape of cVector
   if (!((cShape[0] == mTile * nTile) && (cShape[1] == numElementC)))
     return emitOpError() << "expected matrix C to be shaped (" << mTile * nTile
                          << " x " << numElementC << ")";
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// NVGPU_LdMatrixOp
+//===----------------------------------------------------------------------===//
+LogicalResult LdMatrixOp::verify() {
+
+  // ldmatrix reads data from source in shared memory
+  auto srcMemref = getSrcMemref().getType().cast<MemRefType>();
+
+  // ldmatrix writes data to result/destination in vector registers
+  auto resVector = getRes().getType().cast<VectorType>();
+
+  // vector register shape, element type, and bitwidth
+  ArrayRef<int64_t> resShape = resVector.getShape();
+  Type resType = resVector.getElementType();
+  int64_t elementBitWidth = resType.getIntOrFloatBitWidth();
+
+  // ldmatrix loads 32 bits into vector registers per 8-by-8 tile per thread
+  int64_t numElementsPer32b = 32 / elementBitWidth;
+
+  // number of 8-by-8 tiles
+  int64_t numTiles = getNumTiles();
+
+  // transpose elements in vector registers at 16b granularity when true
+  bool isTranspose = getTranspose();
+
+  // address space id for shared memory
+  unsigned smemAddressSpace = gpu::GPUDialect::getWorkgroupAddressSpace();
+
+  //
+  // verification
+  //
+
+  if (!(srcMemref.getMemorySpaceAsInt() == smemAddressSpace))
+    return emitError()
+           << "expected nvgpu.ldmatrix srcMemref must have memory space "
+           << smemAddressSpace;
+  if (elementBitWidth > 32)
+    return emitError() << "nvgpu.ldmatrix works for 32b or lower";
+  if (isTranspose && !(elementBitWidth == 16))
+    return emitError()
+           << "nvgpu.ldmatrix transpose works only at 16b granularity";
+  if (!(resShape[1] == numElementsPer32b))
+    return emitError() << "expected vector register shape[1] = "
+                       << numElementsPer32b;
+  if (!(resShape[0] == numTiles))
+    return emitError()
+           << "expected vector register shape[0] and numTiles to match";
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"

mlir/test/Dialect/NVGPU/invalid.mlir

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+func.func @ldmatrix_address_space_f16_x4(%arg0: memref<128x128xf16, 2>) ->  vector<4x1xf16> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{expected nvgpu.ldmatrix srcMemref must have memory space 3}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf16, 2> -> vector<4x1xf16>
+  return %a : vector<4x1xf16>
+}
+// -----
+
+func.func @ldmatrix_num_elements_f16_x4(%arg0: memref<128x128xf16, 3>) ->  vector<4x1xf16> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{expected vector register shape[1] = 2}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf16, 3> -> vector<4x1xf16>
+  return %a : vector<4x1xf16>
+}
+// -----
+
+func.func @ldmatrix_num_tiles_f16_x4(%arg0: memref<128x128xf16, 3>) ->  vector<2x2xf16> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{expected vector register shape[0] and numTiles to match}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf16, 3> -> vector<2x2xf16>
+  return %a : vector<2x2xf16>
+}
+// -----
+
+func.func @ldmatrix_num_tiles_f32_x4(%arg0: memref<128x128xf32, 3>) ->  vector<4x2xf32> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{expected vector register shape[1] = 1}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf32, 3> -> vector<4x2xf32>
+  return %a : vector<4x2xf32>
+}
+// -----
+
+func.func @ldmatrix_trans_f32_x4(%arg0: memref<128x128xf32, 3>) ->  vector<4x1xf32> {
+  %c0  = arith.constant 0 : index
+  // expected-error @+1 {{nvgpu.ldmatrix transpose works only at 16b granularity}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = true, numTiles = 4 : i32} : memref<128x128xf32, 3> -> vector<4x1xf32>
+  return %a : vector<4x1xf32>
+}
+// -----
+
+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}}
+  %a = nvgpu.ldmatrix %arg0[%c0, %c0] {transpose = false, numTiles = 4 : i32} : memref<128x128xf32, 3> -> vector<4x2xf16>
+  return %a : vector<4x2xf16>
+}
+// -----
+
 func.func @m16n8k16_fp16_vector_shape_a(%arg0: vector<4x4xf16>, %arg1: vector<2x2xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {
   // expected-error @+1 {{expected 256 warp-wide matrix A elements}}
   %d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<4x4xf16>, vector<2x2xf16>, vector<2x2xf16>) -> vector<2x2xf16>    
   return %d : vector<2x2xf16>
 }
 // -----
 
 func.func @m16n8k16_fp16_vector_shape_b(%arg0: vector<4x2xf16>, %arg1: vector<2x4xf16>, %arg2: vector<2x2xf16>) -> vector<2x2xf16> {