[water] Add ScaledMMAOp

[tgymnich]

• Add ScaledMMAOp
  fixes scaled_mma operation and relevant mma kinds #924

[github-actions]

Water Code Coverage

Filename                                                           Functions  Missed Functions  Executed       Lines      Missed Lines     Cover    Branches   Missed Branches     Cover
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
lib/Transforms/MemrefDecomposition.cpp                                    28                 0   100.00%         600                49    91.83%         104                46    55.77%
lib/Transforms/AllocToAlloca.cpp                                           2                 0   100.00%          17                 0   100.00%           0                 0         -
lib/Transforms/CheckStaticAssertions.cpp                                   2                 0   100.00%          22                 1    95.45%           8                 4    50.00%
lib/Transforms/GPUModuleToBinary.cpp                                      19                 5    73.68%         339               115    66.08%         128                57    55.47%
lib/Transforms/DropTransformOps.cpp                                        2                 0   100.00%          16                 0   100.00%           2                 0   100.00%
lib/Transforms/GPUToGPURuntime.cpp                                        14                 0   100.00%         298                23    92.28%          40                17    57.50%
lib/Transforms/SLPVectorizer.cpp                                          61                 3    95.08%        1065                97    90.89%         558               165    70.43%
lib/Transforms/AccessCheckers.cpp                                         35                 1    97.14%         446                40    91.03%         124                30    75.81%
lib/Transforms/AssembleISA.cpp                                             4                 1    75.00%          30                 2    93.33%           2                 1    50.00%
lib/Dialect/Wave/Transforms/LoweringPatterns.cpp                          48                 2    95.83%         966               146    84.89%         272                82    69.85%
lib/Dialect/Wave/Transforms/PropagateDefaultsFromConstraints.cpp           3                 3     0.00%          35                35     0.00%          12                12     0.00%
lib/Dialect/Wave/Transforms/TypeConverter.cpp                              7                 2    71.43%          96                26    72.92%          32                17    46.88%
lib/Dialect/Wave/Transforms/LowerReadWriteOps.cpp                         10                 0   100.00%         238                18    92.44%          58                11    81.03%
lib/Dialect/Wave/Transforms/DetectNormalForms.cpp                          4                 0   100.00%          48                 0   100.00%           8                 0   100.00%
lib/Dialect/Wave/Transforms/ExpandVariadicReductions.cpp                   2                 0   100.00%          24                 1    95.83%           6                 1    83.33%
lib/Dialect/Wave/Transforms/InferTypes.cpp                               110                14    87.27%        1921               149    92.24%         876               438    50.00%
lib/Dialect/Wave/Transforms/LowerWaveToMLIR.cpp                            5                 0   100.00%         130                 1    99.23%          16                 2    87.50%
lib/Dialect/Wave/Transforms/Utils.cpp                                      6                 0   100.00%          96                 5    94.79%          26                 4    84.62%
lib/Dialect/Wave/Transforms/ResolveDistributedAllocations.cpp              7                 0   100.00%         183                16    91.26%          32                14    56.25%
lib/Dialect/Wave/IR/WaveOps.cpp                                          170                19    88.82%        3132               343    89.05%        1230               269    78.13%
lib/Dialect/Wave/IR/WaveAttrs.cpp                                         73                 6    91.78%         966                97    89.96%         424                62    85.38%
lib/Dialect/Wave/IR/IndexExpr.cpp                                         11                 0   100.00%         119                 1    99.16%          24                 3    87.50%
lib/Dialect/Wave/IR/WaveDialect.cpp                                       14                 0   100.00%         528                20    96.21%         194                13    93.30%
lib/Dialect/Wave/IR/WaveTypes.cpp                                          9                 1    88.89%          75                 8    89.33%          18                 3    83.33%
lib/Dialect/Wave/IR/WaveInterfaces.cpp                                   104                 3    97.12%        1574               100    93.65%         634               104    83.60%
lib/Dialect/Wave/IR/WaveUtils.cpp                                         21                 0   100.00%         190                 8    95.79%          78                13    83.33%
lib/Dialect/NormalForm/Transforms/LowerNormalFormModule.cpp                3                 0   100.00%          34                 6    82.35%           8                 2    75.00%
lib/Dialect/NormalForm/IR/NormalFormDialect.cpp                            1                 0   100.00%           6                 0   100.00%           0                 0         -
lib/Dialect/NormalForm/IR/NormalFormOps.cpp                               12                 0   100.00%         201                 9    95.52%          58                 7    87.93%
lib/Pipelines/Pipelines.cpp                                                2                 0   100.00%          27                 0   100.00%           0                 0         -
lib/Analysis/InUseForSpeculation.cpp                                      12                 1    91.67%         142                 8    94.37%          32                 4    87.50%
include/water/Dialect/Wave/Transforms/LoweringPatterns.h                   1                 0   100.00%           3                 0   100.00%           0                 0         -
include/water/Dialect/Wave/IR/IndexExpr.h                                  1                 0   100.00%          10                 0   100.00%           2                 0   100.00%
include/water/Dialect/Wave/IR/WaveInterfaces.h                            40                 3    92.50%         159                 8    94.97%           8                 2    75.00%
include/water/Dialect/Wave/IR/WaveTypes.h                                  1                 0   100.00%           5                 0   100.00%           4                 0   100.00%
include/water/Dialect/Wave/IR/WaveUtils.h                                  1                 0   100.00%           5                 0   100.00%           4                 1    75.00%
include/water/Dialect/Wave/IR/WaveAttrs.h                                  4                 0   100.00%          14                 0   100.00%           0                 0         -
include/water/Dialect/NormalForm/IR/NormalFormInterfaces.h                 1                 1     0.00%           4                 4     0.00%           0                 0         -
include/water/Analysis/InUseForSpeculation.h                              12                 3    75.00%          39                17    56.41%          16                10    37.50%
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
TOTAL                                                                    862                68    92.11%       13803              1353    90.20%        5038              1394    72.33%

Download full HTML report

[Copilot]

Pull request overview

This PR adds a ScaledMmaOp to the Wave dialect, supporting scaled matrix multiply-accumulate operations using microscaling floating point (MXFP) formats (e.g., FP4 with E8M0 scale factors). It includes the full stack: op definition, type/index inference, lowering to amdgpu.scaled_mfma, Python bindings, and the emitter path.

Changes:

• New ScaledMmaOp in the Wave dialect with 5 operands (lhs, lhs_scale, rhs, rhs_scale, accumulator), full verification, type/index propagation, and elements-per-thread analysis
• Lowering pattern from wave.scaled_mma to amdgpu.scaled_mfma, including scalar extraction for scale factors
• Bug fix in BitcastOp.scale_factor to use get_custom(self.arg).type instead of self.arg.type

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
water/include/water/Dialect/Wave/IR/WaveOps.td	ODS definition of ScaledMmaOp with 5 inputs + kind attribute
water/lib/Dialect/Wave/IR/WaveOps.cpp	Verification, type/index propagation, EPT analysis, and new indexing expressions for F8F6F4 MMA kinds
water/lib/Dialect/Wave/Transforms/LoweringPatterns.cpp	Lowering pattern from ScaledMmaOp to amdgpu.ScaledMFMAOp
wave_lang/kernel/wave/mlir_converter/water_emitter.py	Python emitter support for ScaledMMA node → ScaledMmaOp
wave_lang/kernel/ops/wave_ops.py	_is_fp4 property, MMA_SCALE_FP4 in operand maps, BitcastOp bugfix
lit_tests/kernel/wave/mlir_converter.py	End-to-end test for MXFP4 scaled GEMM
water/test/Dialect/Wave/ops.mlir	Round-trip test for wave.scaled_mma
water/test/Dialect/Wave/python_bindings.py	Python binding tests for ScaledMmaOp
water/test/Dialect/Wave/propagate-elements-per-thread.mlir	EPT backward propagation test for scaled MMA
water/test/Dialect/Wave/lower-wave-to-mlir.mlir	Lowering tests for bitcast and scaled MMA

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

You can also share your feedback on Copilot code review. Take the survey.

[ftynse]

Generally makes sense. My only concern is about isScaled and isFP4 that are not fully documented and the usage seems to be contradictory. I also see pywave has separate flags for LHS and RHS scaling, though I don't know if we have instructions allowing for that.

[ftynse]

Do we care about shared memory in this test? If not, I'd remove to reduce complexity.

[tgymnich]

Yes, shared memory should be i8 and have the right (scaled) shape

[ftynse]

Let's be systematic and either test variable names or not, but not be halfway with only accumulators.

[ftynse]

happy CI is irrelevant