[water] module-level alias for WaveIndexMappingAttr

[ftynse]

Use module-level alias for index mappings to decrease their verbosity in
the overall IR: only the alias name is printed at use sites and
identical mappings have the same alias. This is similar to affine maps.

Signed-off-by: Alex Zinenko git@ozinenko.com

[Copilot]

Pull request overview

This PR introduces module-level attribute aliases for #wave.index_mapping to reduce verbosity in printed Wave IR (similar to affine map aliases), and updates the Wave op assembly format plus test expectations accordingly.

Changes:

• Add an OpAsmDialectInterface to alias wave::WaveIndexMappingAttr with a wim-prefixed module-level alias.
• Switch Wave op index assembly syntax to print/parse the typed attribute directly (e.g., index [#wave.symbol_mapping<...>]), enabling aliasing at use sites.
• Update multiple Wave dialect tests and integration inputs to the new printed form (and add --mlir-print-local-scope where needed).

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
water/lib/Dialect/Wave/IR/WaveDialect.cpp	Registers an ASM dialect interface that aliases WaveIndexMappingAttr as wim*.
water/include/water/Dialect/Wave/IR/WaveOps.td	Updates Wave op assembly format to print/parse $index directly instead of custom<WaveIndexDict>.
water/test/Integration/matmul_lowering.mlir	Updates IR input to use #wim* aliases and #wave.symbol_mapping index form.
water/test/Integration/matmul_lowering_shmem.mlir	Same as above for shared-memory variant.
water/test/Dialect/Wave/resolve-distributed-allocations.mlir	Updates FileCheck expectations to match new index printing form.
water/test/Dialect/Wave/propagate-elements-per-thread.mlir	Updates index attributes in tests to #wave.symbol_mapping<...> form.
water/test/Dialect/Wave/ops.mlir	Updates index attribute spellings in various op syntax tests.
water/test/Dialect/Wave/ops-invalid.mlir	Updates invalid-syntax tests to use the new index attribute form and expected diagnostics.
water/test/Dialect/Wave/lower-wave-to-mlir.mlir	Updates lowering tests to the new index attribute form.
water/test/Dialect/Wave/lower-wave-to-mlir-invalid.mlir	Updates invalid lowering tests to the new index attribute form.
water/test/Dialect/Wave/infer-index-exprs-lattice.mlir	Adds --mlir-print-local-scope and updates checks for the new printed index form.
water/test/Dialect/Wave/expand-variadic-reductions.mlir	Updates reduction expansion checks to the new index form.
water/test/Dialect/Wave/detect-normal-forms.mlir	Updates normal-form detection tests to the new index form.
water/test/Dialect/Wave/attr-type.mlir	Adds --mlir-print-local-scope to stabilize printing with aliases.

[Copilot]

Changing commonArgumentsSyntax from custom<WaveIndexDict>($index) to printing/parsing $index directly changes the textual IR syntax for the index keyword from the legacy index [{M : <...>, ...}] form to index [#wave.symbol_mapping<...>]. If this wasn’t intended as a breaking change for textual Wave IR, consider keeping a custom parser that accepts the old syntax (and/or both syntaxes) while printing the new, alias-friendly form.

Suggested change

	"( `index` $index^ )?"
	"( `index` custom<WaveIndexDict>($index)^ )?"

[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               101    90.52%         558               170    69.53%
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%         963               147    84.74%         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%         219                15    93.15%          54                10    81.48%
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%         129                 1    99.22%          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              4                 0   100.00%          97                13    86.60%          22                10    54.55%
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%        1596               189    88.16%         712               136    80.90%
lib/Dialect/Wave/IR/WaveAttrs.cpp                                         77                 5    93.51%         985                84    91.47%         434                64    85.25%
lib/Dialect/Wave/IR/IndexExpr.cpp                                         10                 0   100.00%         117                 1    99.15%          24                 3    87.50%
lib/Dialect/Wave/IR/WaveDialect.cpp                                       15                 0   100.00%         498                 9    98.19%         174                 6    96.55%
lib/Dialect/Wave/IR/WaveTypes.cpp                                          9                 1    88.89%          75                 8    89.33%          18                 3    83.33%
lib/Dialect/Wave/IR/WaveInterfaces.cpp                                    37                 0   100.00%         661                42    93.65%         322                44    86.34%
lib/Dialect/Wave/IR/WaveUtils.cpp                                         23                 0   100.00%         217                 8    96.31%          84                14    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%          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                                                                    845                69    91.83%       13531              1348    90.04%        5046              1403    72.20%

Download full HTML report

[martin-luecke]

Really cool change to see! The attention kernel is now so much easier to read and comparing whether to indexing expressions are similar is now 1 cycle in my brain.

With this change and not using the custom parser/printer for WaveIndexDict on WaveOp anymore, is this code now dead and could be deleted?

wave/water/include/water/Dialect/Wave/IR/WaveInterfaces.h

Lines 54 to 57 in 5e5fd76

	mlir::ParseResult parseWaveIndexDict(mlir::OpAsmParser &parser,
	mlir::ArrayAttr &out);
	void printWaveIndexDict(mlir::OpAsmPrinter &printer, mlir::Operation *op,
	mlir::ArrayAttr arr);

wave/water/lib/Dialect/Wave/IR/WaveInterfaces.cpp

Lines 235 to 294 in 5e5fd76

	// ODS custom directive: parseWaveIndexDict/printWaveIndexDict
	ParseResult wave::parseWaveIndexDict(OpAsmParser &parser, ArrayAttr &out) {
	auto parseSingleMapping =
	[&](wave::WaveSymbolMappingAttr &out) -> ParseResult {
	SmallVector<std::pair<wave::WaveSymbolAttr, Attribute>> entries;
	MLIRContext *ctx = parser.getContext();
	auto parseEntry = [&]() -> ParseResult {
	StringRef symbolName;
	if (parser.parseKeyword(&symbolName) || parser.parseColon())
	return failure();
	WaveIndexMappingAttr mapping;
	if (failed(parser.parseCustomAttributeWithFallback(mapping)))
	return failure();
	entries.emplace_back(wave::WaveSymbolAttr::get(ctx, symbolName), mapping);
	return success();
	};
	if (parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Braces,
	parseEntry))
	return failure();
	out = wave::WaveSymbolMappingAttr::get(ctx, entries);
	return success();
	};
	SmallVector<Attribute> mappings;
	if (parser.parseCommaSeparatedList(OpAsmParser::Delimiter::Square,
	[&]() -> ParseResult {
	wave::WaveSymbolMappingAttr mapping;
	if (failed(parseSingleMapping(mapping)))
	return failure();
	mappings.push_back(mapping);
	return success();
	}))
	return failure();
	out = parser.getBuilder().getArrayAttr(mappings);
	return success();
	}
	void wave::printWaveIndexDict(OpAsmPrinter &printer, Operation *op,
	ArrayAttr arr) {
	auto printOne = [&](wave::WaveSymbolMappingAttr mapping) {
	printer.getStream() << "{";
	llvm::interleaveComma(
	mapping.getMapping(), printer.getStream(), [&](auto pair) {
	auto [key, value] = pair;
	printer.getStream() << key.getName() << " : ";
	if (auto mappingAttr =
	llvm::dyn_cast<wave::WaveIndexMappingAttr>(value)) {
	mappingAttr.print(printer);
	} else {
	printer.printAttribute(value);
	}
	});
	printer.getStream() << "}";
	};
	printer.getStream() << "[";
	llvm::interleaveComma(arr, printer.getStream(), [&](Attribute a) {
	printOne(llvm::cast<wave::WaveSymbolMappingAttr>(a));
	});
	printer.getStream() << "]";
	}