[water] Add BitcastOp and Scaled Dimensions

[tgymnich]

Introduces wave.bitcast, a new Water dialect op that reinterprets the raw bits of a register-resident tensor as a different element type. Unlike a plain type cast, bitcast preserves the total bit count by scaling exactly one tensor dimension by the source-to-destination element bitwidth ratio.

New dialect pieces

• wave.bitcast op (WaveOps.td, WaveOps.cpp) with verifier, forward/backward type inference, index-expression propagation, and elements-per-thread propagation.
• Scaled dimensions: a dimension is "scaled" when its symbol maps to a WaveExprListAttr ceiling-division expression in the hyperparameters rather than a plain integer. The verifier enforces at most one such dimension per bitcast and validates the ratio against the hyperparameter values.

Hyperparameter extensions

• WaveExprListAttr — a new attribute wrapping a single-result affine map (ceiling division of a symbol by a constant) together with the referenced symbols. Encodes a derived dimension whose value is ceil(base_symbol / divisor).
• WaveHyperparameterAttr::getSymbolValue extended to resolve WaveExprListAttr entries transitively using an iterative worklist (topological-sort style), so derived symbols whose dependencies are also derived resolve correctly.
• verifyOperationAttribute extended to validate WaveExprListAttr entries: single-result, ceiling division of a symbol by a constant, divisor is an integer, dividend is evenly divisible by divisor, no cycles.

Emitter

• Detects sympy expressions of the form base / n (integer ceiling division) and emits a WaveExprListAttr so the Water type system can track the scaled dimension symbolically.

Lowering

• Pattern to lower wave.bitcast between fully-typed wave tensors to an arith.bitcast on the underlying vector.

fixes #922

[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                96    90.99%         558               169    69.71%
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                          46                 2    95.65%         928               144    84.48%         266                80    69.92%
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                               109                14    87.16%        1844               150    91.87%         872               438    49.77%
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                                          151                15    90.07%        2652               272    89.74%        1082               225    79.21%
lib/Dialect/Wave/IR/WaveAttrs.cpp                                         73                 7    90.41%         966               100    89.65%         424                64    84.91%
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                18    96.59%         194                11    94.33%
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               103    83.75%
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                                                                    840                65    92.26%       13208              1281    90.30%        4880              1351    72.32%

Download full HTML report

[Copilot]

Pull request overview

This PR extends the Wave/Water stack to support a new wave.bitcast op and to represent “derived” symbolic dimensions (e.g. K/2) as symbolic hyperparameters via #wave.expr_list, enabling downstream resolution and validation.

Changes:

• Adds BitcastOp end-to-end: Wave op definition, verifier/type+EPT propagation, lowering to vector.bitcast, and Python emitter support.
• Extends #wave.hyperparameters to allow WaveExprListAttr values (derived symbols), including recursive evaluation in getSymbolValue and additional dialect verification.
• Updates Python/CAPI bindings and adds MLIR/Python tests for bitcast + derived hyperparameters (including float4/float8 element types).

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
wave_lang/kernel/wave/mlir_converter/water_emitter.py	Emits wave.bitcast, adds float4/6/8 dtype mapping, and generates derived-dimension symbol names + expr_list hyperparameters.
water/python/WaterExtensionNanobind.cpp	Allows WaveHyperparameterAttr.get() to accept non-int attribute values from Python.
water/lib/CAPI/Dialects.cpp	Permits WaveExprListAttr values in hyperparameter CAPI construction.
water/lib/Dialect/Wave/IR/WaveAttrs.cpp	Implements recursive resolution of derived hyperparameters via WaveExprListAttr.
water/lib/Dialect/Wave/IR/WaveDialect.cpp	Verifies expr_list hyperparameters (single-result maps; referenced symbols must exist).
water/include/water/Dialect/Wave/IR/WaveAttrs.td	Documents hyperparameters supporting integer or expr_list values.
water/include/water/Dialect/Wave/IR/WaveOps.td	Defines BitcastOp in the Wave dialect.
water/lib/Dialect/Wave/IR/WaveOps.cpp	Adds BitcastOp verifier + type/EPT propagation.
water/lib/Dialect/Wave/Transforms/LoweringPatterns.cpp	Lowers wave.bitcast to vector.bitcast.
water/test/Dialect/Wave/python_bindings.py	Adds Python binding tests for expr_list hyperparameters and parsing/printing wave.bitcast.
water/test/Dialect/Wave/ops.mlir	Adds textual IR tests for wave.bitcast.
water/test/Dialect/Wave/infer-types.mlir	Adds type inference tests and negative cases for bitcast dimension scaling.
water/test/Dialect/Wave/propagate-elements-per-thread.mlir	Adds EPT propagation test for bitcast scaling behavior.

---

💡 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 make sense, but we need to iron out the bitcast semantics and tighten verification.

For bitcast, I'm not sure it is the last dimension that is being bitcasted. It may be whatever is the dimension we are (implicitly) vectorizing along, i.e., the dimension where the index expression step is other than 1 (or the last one when all are 1). This will be then consistent with EPT propagation (I don't see any other way) and clearly indicate that index expression propagation is not identity here (it can't be since it must be consistent with EPT propagation). With all that, I'm not sure how exactly to model that during type inference. It appears that type inference is not really possible, unless somehow run in conjunction with index expression inference, since we don't really know which dimension is being bitcast.

[ftynse]

I didn't yet look at python bindings or tests pending design issues, but I notice some missing coverage for bitcast op error messages.

[ftynse]

Reviewed python now. Nothing but nits, though the question of scaled dimension remains unresolved in the op definition.

[ftynse]

Can this land?

[tgymnich]

yes