Add negative tests and remove redundant checks

Author: justinrosner

mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

@@ -451,7 +451,6 @@ LogicalResult ScaledWMMAOp::verify() {
   auto isF8 = llvm::IsaPred<Float8E4M3FNType, Float8E5M2Type>;
   auto isF6 = llvm::IsaPred<Float6E2M3FNType, Float6E3M2FNType>;
   auto isF4 = llvm::IsaPred<Float4E2M1FNType>;
-  auto isSmallFloat = [&](Type t) { return isF4(t) || isF6(t) || isF8(t); };
   auto isScaleF8 = llvm::IsaPred<Float8E8M0FNUType, Float8E4M3FNType>;
   auto isE8M0 = llvm::IsaPred<Float8E8M0FNUType>;
   auto isE4M3 = llvm::IsaPred<Float8E4M3FNType>;
@@ -460,18 +459,10 @@ LogicalResult ScaledWMMAOp::verify() {
   auto sourceBType = cast<VectorType>(getSourceB().getType());
   auto destType = cast<VectorType>(getDestC().getType());
 
-  // Validate output type is F32.
-  if (!destType.getElementType().isF32())
-    return emitOpError("destination must have f32 element type");
-
   // Validate source element types are small floats (fp4/fp6/fp8).
   Type aElemType = sourceAType.getElementType();
   Type bElemType = sourceBType.getElementType();
 
-  if (!isSmallFloat(aElemType) || !isSmallFloat(bElemType))
-    return emitOpError("source operands must have small float element types "
-                       "(fp4/fp6/fp8)");
-
   // Validate vector lengths based on dimensions.
   int64_t m = getM();
   int64_t aLen = sourceAType.getNumElements();

mlir/test/Conversion/AMDGPUToROCDL/wmma-gfx1250.mlir

@@ -164,3 +164,48 @@ func.func @wmma_unsupported_k(%arg0 : vector<8xf16>, %arg1 : vector<8xf32>) {
   amdgpu.wmma 16x16x16 %arg0 * %arg0 + %arg1 : vector<8xf16>, vector<8xf16>, vector<8xf32>
   return
 }
+
+// -----
+
+func.func @scaled_wmma_wrong_output_length(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<16xf32>,
+                                           %arg2 : vector<4xf8E8M0FNU>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op expected output vector of length 8 but got 16}}
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg2[0] * %arg0) * (%arg2[0] * %arg0) + %arg1 : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<16xf32>
+  return
+}
+
+func.func @scaled_wmma_16x16_wrong_sourceA_length(%arg0 : vector<128xf4E2M1FN>, %arg1 : vector<64xf4E2M1FN>,
+                                                  %arg2 : vector<8xf32>, %arg3 : vector<4xf8E8M0FNU>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op for 16x16x128, sourceA must have 64 elements but got 128}}
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E8M0FNU>, vector<128xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<64xf4E2M1FN>, vector<8xf32>
+  return
+}
+
+func.func @scaled_wmma_16x16_wrong_sourceB_length(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<128xf4E2M1FN>,
+                                                  %arg2 : vector<8xf32>, %arg3 : vector<4xf8E8M0FNU>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op for 16x16x128, sourceB must have 64 elements but got 128}}
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E8M0FNU>, vector<128xf4E2M1FN>, vector<8xf32>
+  return
+}
+
+func.func @scaled_wmma_32x16_wrong_sourceA_length(%arg0 : vector<64xf4E2M1FN>, %arg1 : vector<64xf4E2M1FN>,
+                                                  %arg2 : vector<16xf32>, %arg3 : vector<4xf8E4M3FN>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op for 32x16x128, sourceA must have 128 elements but got 64}}
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
+  return
+}
+
+func.func @scaled_wmma_32x16_wrong_sourceB_length(%arg0 : vector<128xf4E2M1FN>, %arg1 : vector<128xf4E2M1FN>,
+                                                  %arg2 : vector<16xf32>, %arg3 : vector<4xf8E4M3FN>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op for 32x16x128, sourceB must have 64 elements but got 128}}
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E4M3FN>, vector<128xf4E2M1FN>, vector<4xf8E4M3FN>, vector<128xf4E2M1FN>, vector<16xf32>
+  return
+}
+
+func.func @scaled_wmma_invalid_type_combination(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<64xf6E2M3FN>,
+                                                %arg2 : vector<8xf32>, %arg3 : vector<4xf8E8M0FNU>,
+                                                %arg4 : vector<4xf8E4M3FN>) {
+  // expected-error@below {{'amdgpu.scaled_wmma' op invalid combination of matrix and scale types}}
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg4[0] * %arg1) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E4M3FN>, vector<64xf6E2M3FN>, vector<8xf32>
+  return
+}