[water] Migrate from DictionaryAttr to WaveIndexExprsAttr for index expressions

This commit uses WaveIndexExprsAttr and WaveIndexEntryAttr to replace
DictionaryAttr for storing index expressions. The key motivation is that
DictionaryAttr alphabetically sorts its entries, but dimension order must
be preserved to match the tensor type's shape order for correct lowering.

Changes:
- Custom parsing/printing that preserves order
- Update IndexExprsLatticeStorage to use WaveIndexExprsAttr
- Update join() logic to use MapVector for order preservation
- Add C API and Python bindings
- Update Python emitter/converter

The join() logic uses LHS-first ordering policy: LHS entries come first
(in LHS order), then new RHS entries (in RHS order). This ensures
deterministic output when joining lattices.

Signed-off-by: tyb0807 <[email protected]>

Author: tyb0807

water/include/water/Dialect/Wave/IR/WaveAttrs.td

@@ -589,4 +589,13 @@ def WaveIndexExprsAttr : AttrDef<WaveDialect, "WaveIndexExprs"> {
   }];
 }
 
+//-----------------------------------------------------------------------------
+// Typed array attributes
+//-----------------------------------------------------------------------------
+
+def WaveIndexExprsArrayAttr : TypedArrayAttrBase<WaveIndexExprsAttr,
+  "array of WaveIndexExprsAttr"> {
+  let constBuilderCall = "$_builder.getArrayAttr($0)";
+}
+
 #endif // WATER_DIALECT_WAVE_WAVEATTRS

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

@@ -380,15 +380,18 @@ class IndexExprsAnalysisInit {
 
 // Lattice for propagating index expressions across wave dialect operations.
 // In addition to the bottom and top states, it can represent a concrete state
-// manifested as a dictionary attribute mapping symbol names to index mappings.
+// manifested as a WaveIndexExprsAttr mapping dimension symbols to index
+// mappings. The entries preserve order, corresponding to the tensor type's
+// dimension order.
+//
 // The JOIN function is defined similarly to other lattices with special
 // handling for combining thread-dependent and thread-independent index
 // expressions.
 class IndexExprsLatticeStorage {
 public:
   IndexExprsLatticeStorage();
   IndexExprsLatticeStorage(const IndexExprsLatticeStorage &value) = default;
-  IndexExprsLatticeStorage(mlir::DictionaryAttr concreteValue);
+  IndexExprsLatticeStorage(wave::WaveIndexExprsAttr concreteValue);
 
   IndexExprsLatticeStorage &
   operator=(const IndexExprsLatticeStorage &other) = default;
@@ -404,15 +407,38 @@ class IndexExprsLatticeStorage {
 
   // Returns the concrete value stored in the lattice instance, be it fully
   // specified or not, or null if the lattice instance is a top or a bottom.
-  mlir::DictionaryAttr getConcreteValue() const;
+  wave::WaveIndexExprsAttr getConcreteValue() const;
 
   // Return the top lattice instance.
   static IndexExprsLatticeStorage top();
 
   // Return the bottom lattice instance.
   static IndexExprsLatticeStorage bottom();
 
-  // Join two lattice instances and return the result.
+  /// Join two lattice instances and return the result.
+  ///
+  /// Ordering semantics:
+  /// - LHS entries come first (in LHS order), then RHS-only entries (in RHS
+  ///   order).
+  /// - Entries with the same dimension have their mappings merged.
+  ///
+  /// Valid usage scenarios:
+  /// 1. **Same dimensions, same order**: Both LHS and RHS have matching
+  ///    dimension order (e.g., both have {M, K}). The result preserves this
+  ///    order. This is the common case for elementwise ops like wave.add.
+  ///
+  /// 2. **MMA ops**: LHS has {M, K}, RHS has {N, K}, accumulator/result has
+  ///    {M, N}. The `ignoredRhsSymbols` parameter filters dimensions that
+  ///    don't apply (e.g., ignore M when propagating from LHS to result).
+  ///    Joins are done incrementally: start with bottom, join LHS (filtered),
+  ///    then join RHS, then join accumulator.
+  ///
+  /// 3. **Iterate ops**: Block arguments are joined with iter_args, and
+  ///    terminator operands with results. Both should have matching tensor
+  ///    types and thus matching dimension order.
+  ///
+  /// If LHS and RHS have conflicting mappings for the same dimension (i.e.,
+  /// mappings that cannot be merged), the result is `top` (conflict).
   static IndexExprsLatticeStorage
   join(const IndexExprsLatticeStorage &lhs, const IndexExprsLatticeStorage &rhs,
        llvm::ArrayRef<mlir::Attribute> ignoredRhsSymbols = {});
@@ -432,13 +458,14 @@ class IndexExprsLatticeStorage {
   // state.
   void unsafeSet(const IndexExprsLatticeStorage &value);
 
-  // Return a new lattice instance with only the provided symbols present.
+  // Return a new lattice instance with only the provided symbols present,
+  // preserving the current order (filtering only, no reordering).
   IndexExprsLatticeStorage
   keepOnlySymbols(llvm::ArrayRef<wave::WaveSymbolAttr> symbols) const;
 
   // Return a new lattice instance where all expressions no longer have
   // references to the provided iterator symbols. Note that this doesn't remove
-  // elements from the mapping dictionary but updates the mapped expressions.
+  // elements from the mapping but updates the mapped expressions.
   IndexExprsLatticeStorage
   withoutIterSymbols(llvm::ArrayRef<wave::WaveSymbolAttr> iterSymbols) const;
 
@@ -447,8 +474,8 @@ class IndexExprsLatticeStorage {
   LLVM_DUMP_METHOD void dump() const;
 
 private:
-  // The internal storage is either a dictionary attribute with one entry per
-  // symbol indexing the value or one of the top/bottom flags.
+  // The internal storage is either a WaveIndexExprsAttr with ordered entries
+  // per dimension symbol, or one of the top/bottom flags.
   llvm::PointerIntPair<mlir::Attribute, 2> value;
 
   // State flags.
@@ -458,6 +485,7 @@ class IndexExprsLatticeStorage {
 };
 
 void operator<<(mlir::Diagnostic &diag, const IndexExprsLatticeStorage &value);
+
 } // namespace wave
 
 namespace llvm {

water/include/water/Dialect/Wave/IR/WaveOps.td

@@ -40,7 +40,7 @@ def WaveMemoryType : Type<Or<[WaveTensorInMemory.predicate, AnyMemRef.predicate]
 class WaveOp<string mnemonic, list<Trait> traits = []>
     : Op<WaveDialect, mnemonic, !listconcat(traits, [HasWaveIndexMapping])> {
   dag commonArguments = (ins
-    Arg<OptionalAttr<DictArrayAttr>, "Index expression">:$index
+    Arg<OptionalAttr<WaveIndexExprsArrayAttr>, "Index expression">:$index
   );
 
   string commonArgumentsSyntax = "( `index` custom<WaveIndexDict>($index)^ )?";

water/include/water/Dialect/Wave/IR/WaveUtils.h

@@ -21,11 +21,11 @@ namespace wave {
 /// Return the position of the dimension that is vectorized based on the index
 /// sequence. The dimension with the largest step is considered to be
 /// vectorized. In case of a tie, take the dimension that is farther in the
-/// index dictionary, which is secretly a list. Return failure when the index
-/// sequence step cannot be evaluated statically.
+/// index expressions. Return failure when the index sequence step cannot be
+/// evaluated statically.
 std::optional<uint64_t>
 getPositionOfVectorizedDim(llvm::ArrayRef<wave::WaveSymbolAttr> shape,
-                           mlir::DictionaryAttr indexDict,
+                           wave::WaveIndexExprsAttr indexExprs,
                            wave::WaveHyperparameterAttr hyper);
 
 // Return the vector shape implied by the index sequence and hyperparameteters,
@@ -34,7 +34,7 @@ getPositionOfVectorizedDim(llvm::ArrayRef<wave::WaveSymbolAttr> shape,
 // it cannot be fully evaluated.
 llvm::SmallVector<int64_t>
 getUncollapsedVectorShape(llvm::ArrayRef<wave::WaveSymbolAttr> shape,
-                          mlir::DictionaryAttr indexDict,
+                          wave::WaveIndexExprsAttr indexExprs,
                           wave::WaveHyperparameterAttr hyper);
 
 /// Resolve named Wave symbols to concrete integer values using the

water/include/water/Dialect/Wave/Transforms/DataFlowAnalyses.h

@@ -18,12 +18,13 @@ class DataFlowSolver;
 class SymbolTableCollection;
 class Operation;
 class Value;
-class DictionaryAttr;
 } // namespace mlir
 
 namespace wave {
+class WaveIndexExprsAttr;
+
 using SetIndexLatticeFn =
-    llvm::function_ref<void(mlir::Value, mlir::DictionaryAttr)>;
+    llvm::function_ref<void(mlir::Value, wave::WaveIndexExprsAttr)>;
 using OverrideInitializationFn = llvm::function_ref<llvm::LogicalResult(
     mlir::Operation *, SetIndexLatticeFn)>;
 

water/include/water/c/Dialects.h

@@ -131,6 +131,51 @@ mlirWaveIndexMappingAttrGetNumSymbols(MlirAttribute attr);
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirWaveIndexMappingAttrGetSymbol(MlirAttribute attr, intptr_t index);
 
+//===---------------------------------------------------------------------===//
+// WaveIndexEntryAttr
+//===---------------------------------------------------------------------===//
+
+/// Checks whether the given MLIR attribute is a WaveIndexEntryAttr.
+MLIR_CAPI_EXPORTED bool mlirAttributeIsAWaveIndexEntryAttr(MlirAttribute attr);
+
+/// Creates a new WaveIndexEntryAttr with the given dimension symbol and
+/// mapping.
+MLIR_CAPI_EXPORTED MlirAttribute mlirWaveIndexEntryAttrGet(
+    MlirContext mlirCtx, MlirAttribute dimension, MlirAttribute mapping);
+
+/// Returns the typeID of a WaveIndexEntryAttr.
+MLIR_CAPI_EXPORTED MlirTypeID mlirWaveIndexEntryAttrGetTypeID();
+
+/// Gets the dimension symbol from a WaveIndexEntryAttr.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirWaveIndexEntryAttrGetDimension(MlirAttribute attr);
+
+/// Gets the mapping from a WaveIndexEntryAttr.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirWaveIndexEntryAttrGetMapping(MlirAttribute attr);
+
+//===---------------------------------------------------------------------===//
+// WaveIndexExprsAttr
+//===---------------------------------------------------------------------===//
+
+/// Checks whether the given MLIR attribute is a WaveIndexExprsAttr.
+MLIR_CAPI_EXPORTED bool mlirAttributeIsAWaveIndexExprsAttr(MlirAttribute attr);
+
+/// Creates a new WaveIndexExprsAttr with the given list of entries.
+MLIR_CAPI_EXPORTED MlirAttribute mlirWaveIndexExprsAttrGet(
+    MlirContext mlirCtx, intptr_t numEntries, MlirAttribute *entries);
+
+/// Returns the typeID of a WaveIndexExprsAttr.
+MLIR_CAPI_EXPORTED MlirTypeID mlirWaveIndexExprsAttrGetTypeID();
+
+/// Gets the number of entries in a WaveIndexExprsAttr.
+MLIR_CAPI_EXPORTED intptr_t
+mlirWaveIndexExprsAttrGetNumEntries(MlirAttribute attr);
+
+/// Gets the entry at the given index from a WaveIndexExprsAttr.
+MLIR_CAPI_EXPORTED MlirAttribute
+mlirWaveIndexExprsAttrGetEntry(MlirAttribute attr, intptr_t index);
+
 //===---------------------------------------------------------------------===//
 // WaveHyperparameterAttr
 //===---------------------------------------------------------------------===//

water/lib/CAPI/Dialects.cpp

@@ -164,6 +164,71 @@ MlirAttribute mlirWaveIndexMappingAttrGetSymbol(MlirAttribute attr,
       llvm::cast<wave::WaveIndexMappingAttr>(unwrap(attr)).getSymbols()[index]);
 }
 
+//===---------------------------------------------------------------------===//
+// WaveIndexEntryAttr
+//===---------------------------------------------------------------------===//
+
+bool mlirAttributeIsAWaveIndexEntryAttr(MlirAttribute attr) {
+  return llvm::isa<wave::WaveIndexEntryAttr>(unwrap(attr));
+}
+
+MlirAttribute mlirWaveIndexEntryAttrGet(MlirContext mlirCtx,
+                                        MlirAttribute dimension,
+                                        MlirAttribute mapping) {
+  MLIRContext *ctx = unwrap(mlirCtx);
+  auto dimAttr = llvm::cast<wave::WaveSymbolAttr>(unwrap(dimension));
+  auto mapAttr = llvm::cast<wave::WaveIndexMappingAttr>(unwrap(mapping));
+  return wrap(wave::WaveIndexEntryAttr::get(ctx, dimAttr, mapAttr));
+}
+
+MlirTypeID mlirWaveIndexEntryAttrGetTypeID() {
+  return wrap(TypeID::get<wave::WaveIndexEntryAttr>());
+}
+
+MlirAttribute mlirWaveIndexEntryAttrGetDimension(MlirAttribute attr) {
+  return wrap(
+      llvm::cast<wave::WaveIndexEntryAttr>(unwrap(attr)).getDimension());
+}
+
+MlirAttribute mlirWaveIndexEntryAttrGetMapping(MlirAttribute attr) {
+  return wrap(llvm::cast<wave::WaveIndexEntryAttr>(unwrap(attr)).getMapping());
+}
+
+//===---------------------------------------------------------------------===//
+// WaveIndexExprsAttr
+//===---------------------------------------------------------------------===//
+
+bool mlirAttributeIsAWaveIndexExprsAttr(MlirAttribute attr) {
+  return llvm::isa<wave::WaveIndexExprsAttr>(unwrap(attr));
+}
+
+MlirAttribute mlirWaveIndexExprsAttrGet(MlirContext mlirCtx,
+                                        intptr_t numEntries,
+                                        MlirAttribute *entries) {
+  MLIRContext *ctx = unwrap(mlirCtx);
+  llvm::SmallVector<wave::WaveIndexEntryAttr> entryAttrs;
+  entryAttrs.reserve(numEntries);
+  for (intptr_t i = 0; i < numEntries; ++i) {
+    entryAttrs.push_back(
+        llvm::cast<wave::WaveIndexEntryAttr>(unwrap(entries[i])));
+  }
+  return wrap(wave::WaveIndexExprsAttr::get(ctx, entryAttrs));
+}
+
+MlirTypeID mlirWaveIndexExprsAttrGetTypeID() {
+  return wrap(TypeID::get<wave::WaveIndexExprsAttr>());
+}
+
+intptr_t mlirWaveIndexExprsAttrGetNumEntries(MlirAttribute attr) {
+  return llvm::cast<wave::WaveIndexExprsAttr>(unwrap(attr)).getEntries().size();
+}
+
+MlirAttribute mlirWaveIndexExprsAttrGetEntry(MlirAttribute attr,
+                                             intptr_t index) {
+  return wrap(
+      llvm::cast<wave::WaveIndexExprsAttr>(unwrap(attr)).getEntries()[index]);
+}
+
 //===---------------------------------------------------------------------===//
 // WaveHyperparameterAttr
 //===---------------------------------------------------------------------===//

water/lib/Dialect/Wave/IR/WaveAttrs.cpp

@@ -722,9 +722,12 @@ Attribute WaveIndexEntryAttr::parse(AsmParser &parser, Type type) {
 
 void WaveIndexEntryAttr::print(AsmPrinter &printer) const {
   // Print: @M : [symbols] -> (start, step, stride)
-  printer.printAttributeWithoutType(getDimension());
+  // Use printStrippedAttrOrType to avoid the #wave.symbol<...> prefix in
+  // generic mode, so we get just <"M"> instead of #wave.symbol<"M">.
+  printer << " ";
+  printer.printStrippedAttrOrType(getDimension());
   printer << " : ";
-  printer.printAttributeWithoutType(getMapping());
+  printer.printStrippedAttrOrType(getMapping());
 }
 
 //===----------------------------------------------------------------------===//
@@ -756,9 +759,10 @@ Attribute WaveIndexExprsAttr::parse(AsmParser &parser, Type type) {
 
 void WaveIndexExprsAttr::print(AsmPrinter &printer) const {
   // Print: <[@M : <mapping>, @K : <mapping>]>
+  // Use printStrippedAttrOrType to avoid #wave<index_entry ...> prefix.
   printer << "<[";
   llvm::interleaveComma(getEntries(), printer, [&](WaveIndexEntryAttr entry) {
-    printer.printAttributeWithoutType(entry);
+    printer.printStrippedAttrOrType(entry);
   });
   printer << "]>";
 }