[mlir][nvgpu] Verify invalid copy size (nfc)

[guraypp]

Summary

This work improves verifier for invalid cases. It is NFC.

[springerm]

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

96

dstElements

98

must

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

 //===- NVGPUDialect.cpp - MLIR NVGPU ops implementation -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the NVGPU dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Verifier.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::nvgpu;
if (dstMemref.getElementType() != srcMemref.getElementType())
     return emitError("source and destination must have the same element type");
   if (size_t(srcMemref.getRank()) != getSrcIndices().size())
     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();
-  if (getBypassL1().has_value()) {
-    int64_t dstElements = getDstElements().getZExtValue();
-    int64_t sizeInBytes =
-        (dstMemref.getElementTypeBitWidth() * dstElements) / 8;
+  int64_t sizeInBytes = (dstMemref.getElementTypeBitWidth() * dstElements) / 8;
+  if (sizeInBytes != 4 && sizeInBytes != 8 && sizeInBytes != 16) {
+    unsigned dstWidth = dstMemref.getElementTypeBitWidth();
+    InFlightDiagnostic diag = emitError();
+    diag << "Requested copy elements is " << dstElements << " with width "
+         << dstMemref.getElementTypeBitWidth()
+         << ". But copy elements could be one of ";
+    if ((32 / dstWidth) > 0)
+      diag << (32 / dstWidth) << ", ";
+    if ((64 / dstWidth) > 0)
+      diag << (64 / dstWidth) << ", ";
+    if ((128 / dstWidth) > 0)
+      diag << (128 / dstWidth) << ".";
+    return diag;
+  }
+  if (getBypassL1().has_value()) {
     int64_t req = 16 * 8 / dstMemref.getElementTypeBitWidth();
     if (getBypassL1().value() && sizeInBytes != 16) {
       return emitOpError() << "bypassL1 does not satify alignment for "
                            << dstMemref << " with destination element "
                            << dstElements
                            << ". Unset bypassL1, or set "
                               "destination element to "
                            << req;

mlir/test/Dialect/NVGPU/invalid.mlir

 // -----
 
 func.func @async_cp_zfill_f32_align1(
   %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index, %srcElements : index) {
     // expected-error @+1 {{'nvgpu.device_async_copy' op bypassL1 does not satify alignment for 'memref<3x16x128xf32, 3>' with destination element 1. Unset bypassL1, or set destination element to 4}}
   %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 1, %srcElements {bypassL1} : memref<128x128xf32> to memref<3x16x128xf32, 3>
   return
 }
+
+// -----
+
+func.func @async_cp_size_invalid_f32(
+  %src: memref<128x128xf32>, %dst: memref<3x16x128xf32, 3>, %i : index) {
+    // expected-error @+1 {{Requested copy elements is 3 with width 32. But copy elements could be one of 1, 2, 4.}}
+  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 3: memref<128x128xf32> to memref<3x16x128xf32, 3>
+  return
+}
+
+// -----
+
+func.func @async_cp_size_invalid_f16(
+  %src: memref<128x128xf16>, %dst: memref<3x16x128xf16, 3>, %i : index) {
+    // expected-error @+1 {{Requested copy elements is 3 with width 16. But copy elements could be one of 2, 4, 8.}}
+  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 3: memref<128x128xf16> to memref<3x16x128xf16, 3>
+  return
+}
+
+// -----
+
+func.func @async_cp_size_invalid_f64(
+  %src: memref<128x128xf64>, %dst: memref<3x16x128xf64, 3>, %i : index) {
+    // expected-error @+1 {{Requested copy elements is 3 with width 64. But copy elements could be one of 1, 2.}}
+  %0 = nvgpu.device_async_copy %src[%i, %i], %dst[%i, %i, %i], 3: memref<128x128xf64> to memref<3x16x128xf64, 3>
+  return
+}

mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir

func.func @optimize_64x16xf32_16x64xf32(%arg0: memref<128x128xf32>,
   %shmB = memref.alloc() : memref<16x64xf32, 3>
 
   // CHECK: [[c6:%.+]] = arith.constant 6 : index
   // CHECK: [[src_bits:%.+]] = arith.andi [[stRow]], [[c6]]
   // CHECK: [[c1:%.+]] = arith.constant 1 : index
   // CHECK: [[xorBits:%.+]] = arith.shli [[src_bits]], [[c1]]
   // CHECK: [[stColPerm:%.+]] = arith.xori [[stCol]], [[xorBits]]
   // CHECK: nvgpu.device_async_copy [[arg0]][[[ldRow]], [[ldCol]]], [[shm]][[[stRow]], [[stColPerm]]]
-  %0 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shm[%stRow, %stCol], 8
+  %0 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shm[%stRow, %stCol], 4
       : memref<128x128xf32> to memref<64x16xf32, 3>
   %1 = nvgpu.device_async_create_group %0
   nvgpu.device_async_wait %1 { numGroups = 1 : i32}
 
   // CHECK: [[c6:%.+]] = arith.constant 6 : index
   // CHECK: [[srcBits:%.+]] = arith.andi [[fragRow]], [[c6]]
   // CHECK: [[c1:%.+]] = arith.constant 1 : index
   // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c1]]
func.func @optimize_64x16xf32_16x64xf32(%arg0: memref<128x128xf32>,
   // Verify 16x64xf32 memory size.
 
   // CHECK: [[c15:%.+]] = arith.constant 15 : index
   // CHECK: [[src_bits:%.+]] = arith.andi [[stRow]], [[c15]]
   // CHECK: [[c2:%.+]] = arith.constant 2 : index
   // CHECK: [[xorBits:%.+]] = arith.shli [[src_bits]], [[c2]]
   // CHECK: [[stColPerm:%.+]] = arith.xori [[stCol]], [[xorBits]]
   // CHECK: nvgpu.device_async_copy [[arg0]][[[ldRow]], [[ldCol]]], [[shmB]][[[stRow]], [[stColPerm]]]
-  %2 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shmB[%stRow, %stCol], 8
+  %2 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shmB[%stRow, %stCol], 4
       : memref<128x128xf32> to memref<16x64xf32, 3>
   %3 = nvgpu.device_async_create_group %0
   nvgpu.device_async_wait %1 { numGroups = 1 : i32}
 
   // CHECK: [[c15:%.+]] = arith.constant 15 : index
   // CHECK: [[srcBits:%.+]] = arith.andi [[fragRow]], [[c15]]
   // CHECK: [[c2:%.+]] = arith.constant 2 : index
   // CHECK: [[xorBits:%.+]] = arith.shli [[srcBits]], [[c2]]
func.func @small_column_size_f64(%arg0: memref<32x32xf64>,
   %shm = memref.alloc() : memref<32x4xf64, 3>
 
   // CHECK: [[c4:%.+]] = arith.constant 4 : index
   // CHECK: [[src_bits:%.+]] = arith.andi [[stRow]], [[c4]]
   // CHECK: [[c1:%.+]] = arith.constant 1 : index
   // CHECK: [[xorBits:%.+]] = arith.shrui [[src_bits]], [[c1]]
   // CHECK: [[stColPerm:%.+]] = arith.xori [[stCol]], [[xorBits]]
   // CHECK: nvgpu.device_async_copy [[arg0]][[[ldRow]], [[ldCol]]], [[shm]][[[stRow]], [[stColPerm]]]
-  %0 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shm[%stRow, %stCol], 8
+  %0 = nvgpu.device_async_copy %arg0[%ldRow, %ldCol], %shm[%stRow, %stCol], 2
       : memref<32x32xf64> to memref<32x4xf64, 3>
   %1 = nvgpu.device_async_create_group %0
   nvgpu.device_async_wait %1 { numGroups = 1 : i32}
 
   // CHECK: [[c6:%.+]] = arith.constant 4 : index
   // CHECK: [[srcBits:%.+]] = arith.andi [[fragRow]], [[c6]]
   // CHECK: [[c1:%.+]] = arith.constant 1 : index
   // CHECK: [[xorBits:%.+]] = arith.shrui [[srcBits]], [[c1]]