[water] Add WaveIndexExprsAttr for ordered index expressions

Add new attributes to represent ordered index expressions that preserve
dimension order, unlike DictionaryAttr which sorts entries alphabetically.

This is the first step toward fixing the index expression ordering issue
where dimension order is lost during lattice operations because:
1. DenseMap loses insertion order during join operations
2. DictionaryAttr::get() always sorts entries alphabetically

New attributes:
- WaveIndexEntryAttr: A single (dimension, mapping) pair
- WaveIndexExprsAttr: An ordered array of entries

The order corresponds to the tensor type's shape dimension order, which
is critical for lowering when the tensor type has been converted to
memref and is no longer available.

Syntax: index_exprs<[@m : <mapping>, @k : <mapping>, @n : <mapping>]>

Signed-off-by: tyb0807 <sontuan.vu@amd.com>

Author: tyb0807

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

@@ -519,4 +519,74 @@ def WaveReadWriteBoundsAttr : AttrDef<WaveDialect, "WaveReadWriteBounds"> {
   }];
 }
 
+//-----------------------------------------------------------------------------
+// Index expression attributes (ordered)
+//-----------------------------------------------------------------------------
+
+def WaveIndexEntryAttr : AttrDef<WaveDialect, "WaveIndexEntry"> {
+  let mnemonic = "index_entry";
+  let description = [{
+    A single entry mapping a tensor dimension symbol to its index mapping.
+
+    This is a component of WaveIndexExprsAttr, representing one dimension's
+    index expression. The dimension is a WaveSymbolAttr (e.g., @M, @K, @N)
+    and the mapping is a WaveIndexMappingAttr specifying start/step/stride.
+
+    Syntax: @M : [symbols] -> (start, step, stride)
+    Example: @M : [#wave.index_symbol<WG0>, #wave.symbol<"BLOCK_M">] -> (WG0 * BLOCK_M, 1, BLOCK_M)
+  }];
+
+  let parameters = (ins
+    "::wave::WaveSymbolAttr":$dimension,
+    "::wave::WaveIndexMappingAttr":$mapping
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+def WaveIndexExprsAttr : AttrDef<WaveDialect, "WaveIndexExprs"> {
+  let mnemonic = "index_exprs";
+  let description = [{
+    An ordered collection of dimension index mappings for Wave tensors.
+
+    Unlike DictionaryAttr which sorts entries alphabetically, this attribute
+    preserves the order of entries as specified. The order corresponds to
+    the dimension order in the associated WaveTensorType's shape.
+
+    Syntax: index_exprs<[@M : <mapping>, @K : <mapping>, @N : <mapping>]>
+
+    The entries are stored in an array, maintaining insertion order.
+    This is critical for lowering where dimension order must match
+    the tensor type's shape, even after the tensor type has been converted
+    to memref.
+  }];
+
+  let parameters = (ins
+    ArrayRefParameter<"::wave::WaveIndexEntryAttr">:$entries
+  );
+
+  let hasCustomAssemblyFormat = 1;
+
+  let extraClassDeclaration = [{
+    /// Look up the index mapping for a given dimension symbol.
+    /// Returns std::nullopt if the dimension is not found.
+    /// Complexity: O(n) where n is the number of dimensions (typically 2-4).
+    std::optional<::wave::WaveIndexMappingAttr>
+    lookup(::wave::WaveSymbolAttr dimension) const;
+
+    /// Look up by dimension name string.
+    std::optional<::wave::WaveIndexMappingAttr>
+    lookup(::llvm::StringRef dimensionName) const;
+
+    /// Get the ordered list of dimension symbols.
+    ::llvm::SmallVector<::wave::WaveSymbolAttr> getDimensions() const;
+
+    /// Get the number of entries.
+    size_t size() const { return getEntries().size(); }
+
+    /// Check if empty.
+    bool empty() const { return getEntries().empty(); }
+  }];
+}
+
 #endif // WATER_DIALECT_WAVE_WAVEATTRS

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

@@ -697,6 +697,98 @@ DeviceConstraintAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// WaveIndexEntryAttr
+//===----------------------------------------------------------------------===//
+
+// Syntax: @M : [symbols] -> (start, step, stride)
+Attribute WaveIndexEntryAttr::parse(AsmParser &parser, Type type) {
+  // Parse dimension symbol: @M
+  WaveSymbolAttr dimension;
+  if (parser.parseCustomAttributeWithFallback<WaveSymbolAttr>(dimension))
+    return {};
+
+  // Parse colon
+  if (parser.parseColon())
+    return {};
+
+  // Parse the index mapping
+  WaveIndexMappingAttr mapping;
+  if (parser.parseCustomAttributeWithFallback<WaveIndexMappingAttr>(mapping))
+    return {};
+
+  return get(parser.getContext(), dimension, mapping);
+}
+
+void WaveIndexEntryAttr::print(AsmPrinter &printer) const {
+  // Print: @M : [symbols] -> (start, step, stride)
+  printer.printAttributeWithoutType(getDimension());
+  printer << " : ";
+  printer.printAttributeWithoutType(getMapping());
+}
+
+//===----------------------------------------------------------------------===//
+// WaveIndexExprsAttr
+//===----------------------------------------------------------------------===//
+
+// Syntax: index_exprs<[@M : <mapping>, @K : <mapping>, @N : <mapping>]>
+Attribute WaveIndexExprsAttr::parse(AsmParser &parser, Type type) {
+  if (parser.parseLess())
+    return {};
+
+  SmallVector<WaveIndexEntryAttr> entries;
+
+  // Parse '[' entries ']' allowing empty or non-empty lists
+  if (parser.parseCommaSeparatedList(AsmParser::Delimiter::Square, [&]() {
+        WaveIndexEntryAttr entry;
+        if (parser.parseCustomAttributeWithFallback<WaveIndexEntryAttr>(entry))
+          return failure();
+        entries.push_back(entry);
+        return success();
+      }))
+    return {};
+
+  if (parser.parseGreater())
+    return {};
+
+  return get(parser.getContext(), entries);
+}
+
+void WaveIndexExprsAttr::print(AsmPrinter &printer) const {
+  // Print: <[@M : <mapping>, @K : <mapping>]>
+  printer << "<[";
+  llvm::interleaveComma(getEntries(), printer, [&](WaveIndexEntryAttr entry) {
+    printer.printAttributeWithoutType(entry);
+  });
+  printer << "]>";
+}
+
+std::optional<WaveIndexMappingAttr>
+WaveIndexExprsAttr::lookup(WaveSymbolAttr dimension) const {
+  for (WaveIndexEntryAttr entry : getEntries()) {
+    if (entry.getDimension() == dimension)
+      return entry.getMapping();
+  }
+  return std::nullopt;
+}
+
+std::optional<WaveIndexMappingAttr>
+WaveIndexExprsAttr::lookup(StringRef dimensionName) const {
+  for (WaveIndexEntryAttr entry : getEntries()) {
+    if (entry.getDimension().getName() == dimensionName)
+      return entry.getMapping();
+  }
+  return std::nullopt;
+}
+
+SmallVector<WaveSymbolAttr> WaveIndexExprsAttr::getDimensions() const {
+  SmallVector<WaveSymbolAttr> dims;
+  dims.reserve(getEntries().size());
+  for (WaveIndexEntryAttr entry : getEntries())
+    dims.push_back(entry.getDimension());
+  return dims;
+}
+
 void wave::WaveDialect::registerAttributes() {
   addAttributes<
 #define GET_ATTRDEF_LIST