Update description of index attributes

Author: justinrosner

mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td

@@ -1248,38 +1248,54 @@ def AMDGPU_ScaledWMMAOp
           VectorOfLengthAndType<[4, 8], [F8E8M0FNU, F8E4M3FN]>:$scaleB,
           ConfinedAttr<I32Attr, [IntIsOneOf<[0, 1]>]>:$scaleBIdx)>,
       Results<(outs ScaledWMMAOutTypes:$destD)> {
+  // TODO: E5M3FNU scales are supported, but there is not yet MLIR support for
+  // this datatype. Once we have support for that, update the scaleA and scaleB
+  // types here.
   let summary = "MLIR wrapper for scaled wmma instructions";
   let description = [{
     The `amdgpu.scaled_wmma` op is an MLIR wrapper around intrinsics for scaled
     `wmma` instructions. These instructions perform matrix multiplication with
     per-block scaling of inputs, supporting fp4, fp6, and fp8 data formats.
 
-    The scale instructions support two tile sizes:
+    The scale instructions support a block size of 16 or 32 and two tile sizes:
     - 16x16x128 with mixed f8/f6/f4 formats (output: vector<4xf32>)
     - 32x16x128 with f4 format only (output: vector<8xf32>)
 
     Scale parameters (`scaleA`, `scaleB`) are small vectors of f8 scale values
-    (either f8E8M0FNU, or f8E4M3FN). The index attributes (`scaleAIdx`, `scaleBIdx`)
-    select which element from the scale vector to use for scaling. During lowering,
-    these vectors are packed into i32/i64 values for the hardware intrinsics.
+    (either f8E8M0FNU, or f8E4M3FN) that are packed into i32/i64 values during
+    lowering. The index attributes (`scaleAIdx`, `scaleBIdx`) select which register
+    lanes provide scale values:
+    - Block size 32: For tile size 16x16x128, each matrix gets 64 scales stored in half
+      a VGPR, with `scaleAIdx`/`scaleBIdx` selecting lanes 0-15 (index=0) or
+      16-31 (index=1). For a tile size of 32x16x128, matrix A gets 128 scales in
+      a full VGPR (`scaleAIdx` is unused), while matrix B gets 64 scales in
+      half a VGPR.
+
+    - Block size 16: For a tile size of 16x16x128, each matrix gets
+      128 scales stored in half of two VGPRs, with `scaleAIdx`/`scaleBIdx`
+      selecting lanes 0-15 (index=0) or 16-31 (index=1) for each of the VGPRs.
+      For 32x16x128, matrix A gets 256 scales in two VGPRs (`scaleAIdx` is unused),
+      while matrix B gets 128 scales stored in half of two VGPRs.
 
     Example:
     ```mlir
       // 16x16x128: fp8 inputs
-      %0 = amdgpu.scaled_wmma 16x16x128 (%scaleVecA[0] * %matA) * (%scaleVecB[0] * %matB) + %matC
+      %0 = amdgpu.scaled_wmma 16x16x128 (%scaleVecA * %matA) * (%scaleVecB * %matB) + %matC
+        {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32}
         : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>,
         vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
 
-      // 32x16x128: fp4 inputs
-      %1 = amdgpu.scaled_wmma 32x16x128 (%scaleVecC[1] * %matD) * (%scaleVecD[0] * %matE) + %matF
+      // 32x16x128: fp4 inputs with different scale indices
+      %1 = amdgpu.scaled_wmma 32x16x128 (%scaleVecD * %matD) * (%scaleVecE * %matE) + %matF
+        {scaleAIdx = 0 : i32, scaleBIdx = 1 : i32}
         : vector<8xf8E4M3FN>, vector<128xf4E2M1FN>,
         vector<8xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
     ```
   }];
   let assemblyFormat = [{
     custom<MNKDimensionList>($m, $n, $k) ` `
-    `(` $scaleA `[` $scaleAIdx `]` `*` $sourceA `)` `*`
-    `(` $scaleB `[` $scaleBIdx `]` `*` $sourceB `)` `+` $destC
+    `(` $scaleA `*` $sourceA `)` `*`
+    `(` $scaleB `*` $sourceB `)` `+` $destC
     attr-dict
     `:` type($scaleA) `,` type($sourceA) `,` type($scaleB) `,` type($sourceB) `,` type($destC)
   }];

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -674,6 +674,8 @@ static Value castScaleOperand(ConversionPatternRewriter &rewriter, Location loc,
   // Handle vector<4xi8> -> i32 or vector<8xi8> -> i64.
   if (auto vectorType = dyn_cast<VectorType>(inputType)) {
     int64_t numElements = vectorType.getNumElements();
+    assert((numElements == 4 || numElements == 8) &&
+           "scale operand must be a vector of length 4 or 8");
     IntegerType outputType =
         (numElements == 4) ? rewriter.getI32Type() : rewriter.getI64Type();
     return LLVM::BitcastOp::create(rewriter, loc, outputType, input);
@@ -691,7 +693,7 @@ static std::optional<uint32_t> getWmmaScaleFormat(Type elemType) {
 }
 
 /// Determines the ROCDL intrinsic name for scaled WMMA based on dimensions
-/// and scale vector length.
+/// and scale block size (16 or 32).
 static std::optional<StringRef>
 getScaledWmmaIntrinsicName(int64_t m, int64_t n, int64_t k, bool isScale16) {
   if (m == 16 && n == 16 && k == 128)

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

@@ -93,10 +93,10 @@ func.func @wmma_k128(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<64xf8E5M2>,
 func.func @wmma_scale_16x16x128_fp8(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<64xf6E2M3FN>,
                                     %arg2 : vector<8xf32>, %arg3 : vector<4xf8E8M0FNU>) {
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg0) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg3 * %arg0) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
 
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtA = 2 : i32, fmtB = 2 : i32, scaleAType = 1 : i32} : (vector<12xi32>, vector<12xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3[1] * %arg1) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<8xf32>
+  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg1) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 1 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<8xf32>
 
   func.return
 }
@@ -105,10 +105,10 @@ func.func @wmma_scale_16x16x128_fp8(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<
 func.func @wmma_scale_16x16x128_fp6(%arg0 : vector<64xf6E2M3FN>, %arg1 : vector<64xf6E3M2FN>,
                                     %arg2 : vector<8xf32>, %arg3 : vector<4xf8E8M0FNU>) {
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtA = 2 : i32, fmtB = 2 : i32} : (vector<12xi32>, vector<12xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg0) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<8xf32>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg3 * %arg0) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<8xf32>
 
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtA = 3 : i32, fmtB = 3 : i32} : (vector<12xi32>, vector<12xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg1) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<4xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf32>
+  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg1) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<4xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf32>
 
   func.return
 }
@@ -118,10 +118,10 @@ func.func @wmma_scale_16x16x128_mixed(%arg0 : vector<64xf8E4M3FN>, %arg1 : vecto
                                       %arg2 : vector<64xf4E2M1FN>, %arg3 : vector<8xf32>,
                                       %arg4 : vector<4xf8E8M0FNU>, %arg5 : vector<4xf8E4M3FN>) {
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg3, {{.*}}, {{.*}} {fmtB = 4 : i32, fmtScaleB = 2 : i32} : (vector<16xi32>, vector<8xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %0 = amdgpu.scaled_wmma 16x16x128 (%arg4[0] * %arg0) * (%arg5[0] * %arg2) + %arg3 : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<8xf32>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg4 * %arg0) * (%arg5 * %arg2) + %arg3 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<8xf32>
 
   // CHECK: rocdl.wmma.scale.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg3, {{.*}}, {{.*}} {fmtA = 2 : i32, fmtB = 4 : i32, fmtScaleB = 2 : i32} : (vector<12xi32>, vector<8xi32>, vector<8xf32>, i32, i32) -> vector<8xf32>
-  %1 = amdgpu.scaled_wmma 16x16x128 (%arg4[0] * %arg1) * (%arg5[0] * %arg2) + %arg3 : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<8xf32>
+  %1 = amdgpu.scaled_wmma 16x16x128 (%arg4 * %arg1) * (%arg5 * %arg2) + %arg3 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf6E2M3FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<8xf32>
 
   func.return
 }
@@ -130,10 +130,10 @@ func.func @wmma_scale_16x16x128_mixed(%arg0 : vector<64xf8E4M3FN>, %arg1 : vecto
 func.func @wmma_scale16_16x16x128_fp8(%arg0 : vector<64xf8E4M3FN>, %arg1 : vector<64xf6E3M2FN>,
                                       %arg2 : vector<8xf32>, %arg3 : vector<8xf8E8M0FNU>) {
   // CHECK: rocdl.wmma.scale16.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} : (vector<16xi32>, vector<16xi32>, vector<8xf32>, i64, i64) -> vector<8xf32>
-  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg0) + %arg2 : vector<8xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg3 * %arg0) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<8xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<8xf32>
 
   // CHECK: rocdl.wmma.scale16.f32.16x16x128.f8f6f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtA = 3 : i32, fmtB = 3 : i32, scaleAType = 1 : i32} : (vector<12xi32>, vector<12xi32>, vector<8xf32>, i64, i64) -> vector<8xf32>
-  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3[1] * %arg1) * (%arg3[0] * %arg1) + %arg2 : vector<8xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf32>
+  %1 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg1) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 1 : i32, scaleBIdx = 0 : i32} : vector<8xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf8E8M0FNU>, vector<64xf6E3M2FN>, vector<8xf32>
 
   func.return
 }
@@ -142,7 +142,7 @@ func.func @wmma_scale16_16x16x128_fp8(%arg0 : vector<64xf8E4M3FN>, %arg1 : vecto
 func.func @wmma_scale_32x16x128_fp4(%arg0 : vector<128xf4E2M1FN>, %arg1 : vector<64xf4E2M1FN>,
                                     %arg2 : vector<16xf32>, %arg3 : vector<4xf8E4M3FN>) {
   // CHECK: rocdl.wmma.scale.f32.32x16x128.f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtScaleA = 2 : i32, fmtScaleB = 2 : i32} : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i32, i32) -> vector<16xf32>
-  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<4xf8E4M3FN>, vector<128xf4E2M1FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E4M3FN>, vector<128xf4E2M1FN>, vector<4xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
 
   func.return
 }
@@ -151,7 +151,7 @@ func.func @wmma_scale_32x16x128_fp4(%arg0 : vector<128xf4E2M1FN>, %arg1 : vector
 func.func @wmma_scale16_32x16x128_fp4(%arg0 : vector<128xf4E2M1FN>, %arg1 : vector<64xf4E2M1FN>,
                                       %arg2 : vector<16xf32>, %arg3 : vector<8xf8E4M3FN>) {
   // CHECK: rocdl.wmma.scale16.f32.32x16x128.f4 {{.*}}, {{.*}}, %arg2, {{.*}}, {{.*}} {fmtScaleA = 2 : i32, fmtScaleB = 2 : i32} : (vector<16xi32>, vector<8xi32>, vector<16xf32>, i64, i64) -> vector<16xf32>
-  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3[0] * %arg0) * (%arg3[0] * %arg1) + %arg2 : vector<8xf8E4M3FN>, vector<128xf4E2M1FN>, vector<8xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<8xf8E4M3FN>, vector<128xf4E2M1FN>, vector<8xf8E4M3FN>, vector<64xf4E2M1FN>, vector<16xf32>
 
   func.return
 }
@@ -170,42 +170,42 @@ func.func @wmma_unsupported_k(%arg0 : vector<8xf16>, %arg1 : vector<8xf32>) {
 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>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg2 * %arg0) * (%arg2 * %arg0) + %arg1 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : 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>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : 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>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : 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>
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : 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>
+  %0 = amdgpu.scaled_wmma 32x16x128 (%arg3 * %arg0) * (%arg3 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : 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>
+  %0 = amdgpu.scaled_wmma 16x16x128 (%arg3 * %arg0) * (%arg4 * %arg1) + %arg2 {scaleAIdx = 0 : i32, scaleBIdx = 0 : i32} : vector<4xf8E8M0FNU>, vector<64xf8E4M3FN>, vector<4xf8E4M3FN>, vector<64xf6E2M3FN>, vector<8xf32>
   return
 }