[Mlir-commits] [mlir] 7131244 - [mlir][AMDGPU] Allow packing of exactly 4 elements. (#181843)

[llvmlistbot at llvm.org]

Author: Muzammiluddin Syed
Date: 2026-02-17T12:01:18-05:00
New Revision: 7131244f383161b000b36da60df63dfbb22b3a52

URL: https://github.com/llvm/llvm-project/commit/7131244f383161b000b36da60df63dfbb22b3a52
DIFF: https://github.com/llvm/llvm-project/commit/7131244f383161b000b36da60df63dfbb22b3a52.diff

LOG: [mlir][AMDGPU] Allow packing of exactly 4 elements. (#181843)

`amdgpu.scaled_mfma` ops ingest byte sized scales stored in 4-byte
registers. To avoid unnecessary padding (where we only ever use the
first byte in this 4-byte register), this canonicalization finds
opportunities to enable packing multiple scales into 4-byte chunks
whenever possible. Note this is necessary but not sufficient to avoid
byte loads from LDS.

This canonicalization should try to pack scales that are extracted from
an alloc in shared mem of size 4 bytes or larger (meaning packing to 4
bytes is possible). Currently we bail out if it is exactly 4 bytes long
which is incorrect and fixed in this PR.

---------

Signed-off-by: Muzammiluddin Syed <muzasyed at amd.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
    mlir/test/Dialect/AMDGPU/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
index f9f11c1f9e540..f452d2de15dc8 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
@@ -1145,9 +1145,9 @@ struct PackScales final : OpRewritePattern<ScaledMFMAOp> {
       }
 
       int64_t numElements = scaleSrcType.getNumElements();
-      if (numElements <= 4) {
+      if (numElements < 4) {
         return rewriter.notifyMatchFailure(
-            op, "no packing if # of scales less than four");
+            op, "do not pack if # of scales less than four");
       }
 
       // Find a linearized idx using the size and offsets of the extract op.

diff  --git a/mlir/test/Dialect/AMDGPU/canonicalize.mlir b/mlir/test/Dialect/AMDGPU/canonicalize.mlir
index 052fabfb300bb..f51d0d0d12ed0 100644
--- a/mlir/test/Dialect/AMDGPU/canonicalize.mlir
+++ b/mlir/test/Dialect/AMDGPU/canonicalize.mlir
@@ -206,6 +206,21 @@ func.func @scaled_mfma_less_than_4(%opA: vector<32xf4E2M1FN>, %opB: vector<32xf4
 
 // -----
 
+// CHECK-LABEL: func @scaled_mfma_exactly_4
+// CHECK: amdgpu.scaled_mfma 16x16x128 ({{.*}}[0] * {{.*}}) * ({{.*}}[1] * {{.*}}
+func.func @scaled_mfma_exactly_4(%opA: vector<32xf4E2M1FN>, %opB: vector<32xf4E2M1FN>, %scalesA: vector<4xf8E8M0FNU>, %scalesB: vector<4xf8E8M0FNU>) -> vector<4xf32> {
+  %cst_0 = arith.constant dense<0.000000e+00> : vector<4xf32>
+  %cst_1 = arith.constant dense<5.877470e-39> : vector<4xf8E8M0FNU>
+  %scaleA = vector.extract %scalesA[0] : f8E8M0FNU from vector<4xf8E8M0FNU>
+  %sA = vector.insert %scaleA, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>
+  %scaleB = vector.extract %scalesB[1] : f8E8M0FNU from vector<4xf8E8M0FNU>
+  %sB = vector.insert %scaleB, %cst_1 [0] : f8E8M0FNU into vector<4xf8E8M0FNU>
+  %res_0 = amdgpu.scaled_mfma 16x16x128 (%sA[0] * %opA) * (%sB[0] * %opB) + %cst_0 : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>
+  return %res_0 : vector<4xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @scaled_mfma_ugly_shapes
 // CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[0] * %{{.*}}) * (%{{.*}}[3] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>
 // CHECK: amdgpu.scaled_mfma 16x16x128 (%{{.*}}[1] * %{{.*}}) * (%{{.*}}[3] * %arg1) + %cst : vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf8E8M0FNU>, vector<32xf4E2M1FN>, vector<4xf32>