[water] Add ordered_syms attribute for memref read/write lowering

Add ordered_syms attribute to wave.read and wave.write operations to
preserve dimension ordering information when the memory operand is a
MemRefType. This is needed because the index dictionary keys are
dimension names, and we need to know their ordering for correct lowering.

The ResolveDistributedAllocations pass now sets this attribute on
read/write ops that use resolved allocate results, extracting the
ordered symbols from the original WaveTensorType before it becomes
a MemRefType.

The LowerReadWriteOps code now handles both WaveTensorType (using
getShape()) and MemRefType (using ordered_syms attribute) memory
operands.

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

Author: tyb0807

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

@@ -310,14 +310,21 @@ def ReadOp : WaveOp<"read", [
   let description = [{
     Moves data from a memory-resident tensor to a register-resident tensor
     preserving the shape.
+
+    When the memory operand is a MemRefType (after ResolveDistributedAllocations),
+    the `ordered_syms` attribute provides the ordered dimension symbols that
+    correspond to the memref dimensions. This is needed because the index
+    dictionary keys are dimension names, and we need to know their ordering.
   }];
 
   let arguments = !con((ins
     Arg<WaveMemoryType, "Memory to read from">:$memory,
     Arg<OptionalAttr<I64Attr>,
         "Number of elements processed by each thread">:$elements_per_thread,
     Arg<OptionalAttr<WaveReadWriteBoundsAttr>,
-        "Bound expressions for each symbolic dimension">:$bounds
+        "Bound expressions for each symbolic dimension">:$bounds,
+    Arg<OptionalAttr<ArrayAttr>,
+        "Ordered dimension symbols for memref operands">:$ordered_syms
   ), commonArguments);
 
   let results = (outs
@@ -364,6 +371,11 @@ def WriteOp : WaveOp<"write", [
   let description = [{
     Moves data from a register-resident tensor into a memory-resident tensor
     preserving the shape.
+
+    When the memory operand is a MemRefType (after ResolveDistributedAllocations),
+    the `ordered_syms` attribute provides the ordered dimension symbols that
+    correspond to the memref dimensions. This is needed because the index
+    dictionary keys are dimension names, and we need to know their ordering.
   }];
 
   let arguments = !con((ins
@@ -372,7 +384,9 @@ def WriteOp : WaveOp<"write", [
     Arg<OptionalAttr<I64Attr>,
         "Number of elements processed by each thread">:$elements_per_thread,
     Arg<OptionalAttr<WaveReadWriteBoundsAttr>,
-        "Bound expressions for each symbolic dimension">:$bounds
+        "Bound expressions for each symbolic dimension">:$bounds,
+    Arg<OptionalAttr<ArrayAttr>,
+        "Ordered dimension symbols for memref operands">:$ordered_syms
   ), commonArguments);
 
   let assemblyFormat =

water/lib/Dialect/Wave/Transforms/LowerReadWriteOps.cpp

@@ -432,14 +432,28 @@ static FailureOr<MemAccessInfo>
 createMemoryIndicesAndMask(ConversionPatternRewriter &rewriter,
                            const TypeConverter *typeConverter, OpTy op,
                            Type memoryTypeArg, VectorType vectorType) {
-  auto memoryType = dyn_cast<wave::WaveTensorType>(memoryTypeArg);
-  if (!memoryType)
-    return rewriter.notifyMatchFailure(
-        op, "lowering with MemRefType memory not yet implemented");
-
   int64_t elementsPerThread = vectorType.getNumElements();
 
-  ArrayRef<wave::WaveSymbolAttr> orderedSyms = memoryType.getShape();
+  // Get ordered symbols from either WaveTensorType or ordered_syms attribute.
+  SmallVector<wave::WaveSymbolAttr> orderedSymsStorage;
+  ArrayRef<wave::WaveSymbolAttr> orderedSyms;
+
+  if (auto memoryType = dyn_cast<wave::WaveTensorType>(memoryTypeArg)) {
+    orderedSyms = memoryType.getShape();
+  } else if (isa<MemRefType>(memoryTypeArg)) {
+    // For MemRefType, get ordered symbols from the ordered_syms attribute.
+    ArrayAttr orderedSymsAttr = op.getOrderedSymsAttr();
+    if (!orderedSymsAttr)
+      return rewriter.notifyMatchFailure(
+          op, "MemRefType memory requires ordered_syms attribute");
+    orderedSymsStorage.reserve(orderedSymsAttr.size());
+    for (Attribute attr : orderedSymsAttr)
+      orderedSymsStorage.push_back(cast<wave::WaveSymbolAttr>(attr));
+    orderedSyms = orderedSymsStorage;
+  } else {
+    return rewriter.notifyMatchFailure(
+        op, "unsupported memory type for lowering");
+  }
 
   wave::WaveReadWriteBoundsAttr boundsDict = op.getBoundsAttr();
   wave::WaveHyperparameterAttr hyper =

water/lib/Dialect/Wave/Transforms/ResolveDistributedAllocations.cpp

@@ -55,6 +55,23 @@ struct ResolveDistributedAllocations
         return;
       }
 
+      // Extract ordered dimension symbols from the original tensor type.
+      ArrayRef<WaveSymbolAttr> shapeSyms = tensorType.getShape();
+      SmallVector<Attribute> orderedSyms(shapeSyms.begin(), shapeSyms.end());
+      ArrayAttr orderedSymsAttr =
+          ArrayAttr::get(allocateOp.getContext(), orderedSyms);
+
+      // Update read/write ops that use this allocate result to include
+      // ordered_syms, since they will need it for lowering after the memory
+      // operand becomes a MemRefType.
+      for (Operation *user : allocateOp.getResult().getUsers()) {
+        if (auto readOp = dyn_cast<ReadOp>(user))
+          readOp.setOrderedSymsAttr(orderedSymsAttr);
+        else if (auto writeOp = dyn_cast<WriteOp>(user))
+          writeOp.setOrderedSymsAttr(orderedSymsAttr);
+      }
+
+      // Update the result type in place.
       allocateOp.getResult().setType(memrefType);
     });
     return result;

water/test/Dialect/Wave/lower-wave-to-mlir.mlir

@@ -867,3 +867,31 @@ module attributes {wave.normal_form = #wave.normal_form<full_types,memory_only_t
     return
   }
 }
+
+// -----
+
+// Test read/write lowering with MemRefType memory operand.
+// This simulates the state after ResolveDistributedAllocations pass has run.
+module attributes {wave.normal_form = #wave.normal_form<full_types,memory_only_types,resolved_allocations>} {
+  // CHECK-LABEL: @lower_read_write_memref
+  func.func @lower_read_write_memref(%mem: memref<64x64xf16, #gpu.address_space<workgroup>>)
+      attributes {wave.hyperparameters = #wave.hyperparameters<{BLOCK_M = 64, BLOCK_N = 64}>} {
+    // Test ReadOp lowering when memory operand is already a MemRefType
+    // and ordered_syms attribute provides the dimension ordering.
+    // CHECK: %[[READ:.*]] = vector.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x64xf16, #gpu.address_space<workgroup>>, vector<8xf16>
+    %0 = wave.read %mem index [{
+        BLOCK_M : [#wave.index_symbol<T0>, #wave.symbol<"BLOCK_M">] -> (T0 mod 64, 1, 64),
+        BLOCK_N : [#wave.index_symbol<T1>, #wave.symbol<"BLOCK_N">] -> (T1 * 8, 8, 1)
+      }] { ordered_syms = [#wave.symbol<"BLOCK_M">, #wave.symbol<"BLOCK_N">] }
+      : (memref<64x64xf16, #gpu.address_space<workgroup>>) -> vector<8xf16>
+
+    // Test WriteOp lowering when memory operand is already a MemRefType.
+    // CHECK: vector.store %[[READ]], %{{.*}}[%{{.*}}, %{{.*}}] : memref<64x64xf16, #gpu.address_space<workgroup>>, vector<8xf16>
+    wave.write %0, %mem index [{
+        BLOCK_M : [#wave.index_symbol<T0>, #wave.symbol<"BLOCK_M">] -> (T0 mod 64, 1, 64),
+        BLOCK_N : [#wave.index_symbol<T1>, #wave.symbol<"BLOCK_N">] -> (T1 * 8, 8, 1)
+      }] { ordered_syms = [#wave.symbol<"BLOCK_M">, #wave.symbol<"BLOCK_N">] }
+      : vector<8xf16>, memref<64x64xf16, #gpu.address_space<workgroup>>
+    return
+  }
+}