[wave2water] Add MXFP4 support to fx_emitter and fix roundtrip comparison

[martin-luecke]

Teach the MLIR-to-FX importer to handle ops used by the MXFP4 GEMM kernel and fix the roundtrip comparison for MMA index expressions.

FX emitter:

• Add _handle_mma_like_op handling both wave.mma and wave.scaled_mma
• Add _handle_bitcast_op for wave.bitcast.
• Add derived dimension support: hyperparameter parsing separates base integers from WaveExprListAttr expressions (e.g. K32 = ceiling(K/32)), with _resolve_dim mapping MLIR dimension names back to symbolic expressions.
• Thread parse_ctx through all _convert_supported_attrs call sites.

Attribute converter:

• Replace convert_index_mapping_array_to_sympy (which dispatched on op.name strings) with convert_index_mapping_dict_to_sympy (single mapping) and convert_mma_index_to_sympy (MMA/ScaledMMA reconstruction with derived_dims support).
• Assert lhs_scale/rhs_scale index consistency for ScaledMMA.

Roundtrip comparison:

• Compare MMA/ScaledMMA indices per-operand (lhs_index, rhs_index, etc.) rather than on the combined dict. The combined dict carries Piecewise expressions over operand-selector symbols that are dead for certain dtype configurations and lost during MLIR serialization.
• Add WaveConstraint comparator and thread subs into assert_traces_equivalent

[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               163    70.79%
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%         965               147    84.77%         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%         236                18    92.37%          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                                62                11    82.26%         920                94    89.78%         362               153    57.73%
lib/Dialect/Wave/Transforms/LowerWaveToMLIR.cpp                            5                 0   100.00%         130                 1    99.23%          16                 2    87.50%
lib/Dialect/Wave/Transforms/InferIndexExprs.cpp                            3                 0   100.00%          34                 1    97.06%           8                 1    87.50%
lib/Dialect/Wave/Transforms/Utils.cpp                                      4                 0   100.00%          64                 5    92.19%          20                 4    80.00%
lib/Dialect/Wave/Transforms/ResolveDistributedAllocations.cpp              7                 0   100.00%         183                16    91.26%          32                14    56.25%
lib/Dialect/Wave/IR/IndexExprInference.cpp                               173                17    90.17%        3269               286    91.25%        1368               490    64.18%
lib/Dialect/Wave/IR/WaveOps.cpp                                           91                10    89.01%        1645               186    88.69%         742               145    80.46%
lib/Dialect/Wave/IR/WaveAttrs.cpp                                         73                 5    93.15%         948                83    91.24%         418                59    85.89%
lib/Dialect/Wave/IR/IndexExpr.cpp                                         10                 0   100.00%         117                 1    99.15%          24                 3    87.50%
lib/Dialect/Wave/IR/WaveDialect.cpp                                       14                 0   100.00%         485                 9    98.14%         170                 5    97.06%
lib/Dialect/Wave/IR/WaveTypes.cpp                                          9                 1    88.89%          75                 8    89.33%          18                 3    83.33%
lib/Dialect/Wave/IR/WaveInterfaces.cpp                                    45                 0   100.00%         781                47    93.98%         338                51    84.91%
lib/Dialect/Wave/IR/WaveUtils.cpp                                         23                 0   100.00%         220                 8    96.36%          84                13    84.52%
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%          16                 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                                                                    851                69    91.89%       13759              1351    90.18%        5086              1410    72.28%

Download full HTML report

[ftynse]

LGTM % an ordering issue with index attrs.