[water] Add lowering pattern using LDS promotion for write group

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

Author: tyb0807

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

@@ -306,7 +306,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<WaveMemoryAccessPatternAttr>,
+        "Memory access pattern controlling LDS promotion">:$memory_access_pattern
   ), commonArguments);
 
   let assemblyFormat =

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

@@ -460,24 +460,24 @@ DeviceConstraintAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 // WaveMemoryAccessPatternAttr
 //===----------------------------------------------------------------------===//
 
-LogicalResult
-WaveMemoryAccessPatternAttr::verify(function_ref<InFlightDiagnostic()> emitError,
-                                    bool use_lds_promotion, StringRef group_id,
-                                    WaveExprListAttr lds_block_global_base,
-                                    WaveExprListAttr lds_block_shape,
-                                    WaveExprListAttr lds_load_indices,
-                                    WaveExprListAttr lds_load_vector_sizes,
-                                    WaveExprListAttr global_store_indices) {
+LogicalResult WaveMemoryAccessPatternAttr::verify(
+    function_ref<InFlightDiagnostic()> emitError, bool use_lds_promotion,
+    StringRef group_id, WaveExprListAttr lds_block_global_base,
+    WaveExprListAttr lds_block_shape, WaveExprListAttr lds_load_indices,
+    WaveExprListAttr lds_load_vector_sizes,
+    WaveExprListAttr global_store_indices) {
   // Validate group_id is not empty
   if (group_id.empty()) {
     return emitError() << "group_id cannot be empty";
   }
 
-  // When LDS promotion is disabled, no LDS-related parameters should be specified
+  // When LDS promotion is disabled, no LDS-related parameters should be
+  // specified
   if (!use_lds_promotion) {
     if (lds_block_global_base || lds_block_shape || lds_load_indices ||
         lds_load_vector_sizes || global_store_indices) {
-      return emitError() << "LDS promotion parameters should not be specified when use_lds_promotion=false";
+      return emitError() << "LDS promotion parameters should not be specified "
+                            "when use_lds_promotion=false";
     }
     return success();
   }
@@ -492,10 +492,12 @@ WaveMemoryAccessPatternAttr::verify(function_ref<InFlightDiagnostic()> emitError
   // Check for partial specification - either all or none should be provided
   if (hasLdsBase || hasLdsShape || hasLdsLoadIndices || hasLdsLoadVectorSizes ||
       hasGlobalStoreIndices) {
-    if (!hasLdsBase || !hasLdsShape || !hasLdsLoadIndices || !hasLdsLoadVectorSizes ||
-        !hasGlobalStoreIndices) {
-      return emitError() << "when LDS promotion is enabled, all LDS parameters must be specified: "
-                            "lds_block_global_base, lds_block_shape, lds_load_indices, lds_load_vector_sizes, "
+    if (!hasLdsBase || !hasLdsShape || !hasLdsLoadIndices ||
+        !hasLdsLoadVectorSizes || !hasGlobalStoreIndices) {
+      return emitError() << "when LDS promotion is enabled, all LDS parameters "
+                            "must be specified: "
+                            "lds_block_global_base, lds_block_shape, "
+                            "lds_load_indices, lds_load_vector_sizes, "
                             "global_store_indices";
     }
   }
@@ -504,13 +506,15 @@ WaveMemoryAccessPatternAttr::verify(function_ref<InFlightDiagnostic()> emitError
   if (hasLdsBase && hasLdsShape && hasLdsLoadIndices && hasLdsLoadVectorSizes &&
       hasGlobalStoreIndices) {
 
-    // Validate that lds_block_global_base and lds_block_shape have consistent ranks
+    // Validate that lds_block_global_base and lds_block_shape have consistent
+    // ranks
     unsigned ldsBaseRank = lds_block_global_base.getRank();
     unsigned ldsShapeRank = lds_block_shape.getRank();
 
     if (ldsBaseRank != ldsShapeRank) {
       return emitError() << "lds_block_global_base rank (" << ldsBaseRank
-                         << ") must match lds_block_shape rank (" << ldsShapeRank << ")";
+                         << ") must match lds_block_shape rank ("
+                         << ldsShapeRank << ")";
     }
 
     // Validate that load indices and vector sizes have consistent ranks
@@ -520,58 +524,68 @@ WaveMemoryAccessPatternAttr::verify(function_ref<InFlightDiagnostic()> emitError
 
     if (ldsLoadIndicesRank != ldsLoadVectorSizesRank) {
       return emitError() << "lds_load_indices rank (" << ldsLoadIndicesRank
-                         << ") must match lds_load_vector_sizes rank (" << ldsLoadVectorSizesRank << ")";
+                         << ") must match lds_load_vector_sizes rank ("
+                         << ldsLoadVectorSizesRank << ")";
     }
 
     if (ldsBaseRank != ldsLoadIndicesRank) {
       return emitError() << "LDS block rank (" << ldsBaseRank
-                         << ") must match LDS load indices rank (" << ldsLoadIndicesRank << ")";
+                         << ") must match LDS load indices rank ("
+                         << ldsLoadIndicesRank << ")";
     }
 
     if (ldsBaseRank != globalStoreIndicesRank) {
       return emitError() << "LDS block rank (" << ldsBaseRank
-                         << ") must match global store indices rank (" << globalStoreIndicesRank << ")";
+                         << ") must match global store indices rank ("
+                         << globalStoreIndicesRank << ")";
     }
 
     // Validate that all symbols are WaveSymbolAttr or WaveIndexSymbolAttr
     if (!llvm::all_of(lds_block_global_base.getSymbols(),
                       llvm::IsaPred<WaveSymbolAttr, WaveIndexSymbolAttr>)) {
-      return emitError() << "lds_block_global_base must only contain WaveSymbolAttr or WaveIndexSymbolAttr";
+      return emitError() << "lds_block_global_base must only contain "
+                            "WaveSymbolAttr or WaveIndexSymbolAttr";
     }
 
     if (!llvm::all_of(lds_block_shape.getSymbols(),
                       llvm::IsaPred<WaveSymbolAttr, WaveIndexSymbolAttr>)) {
-      return emitError() << "lds_block_shape must only contain WaveSymbolAttr or WaveIndexSymbolAttr";
+      return emitError() << "lds_block_shape must only contain WaveSymbolAttr "
+                            "or WaveIndexSymbolAttr";
     }
 
     if (!llvm::all_of(lds_load_indices.getSymbols(),
                       llvm::IsaPred<WaveSymbolAttr, WaveIndexSymbolAttr>)) {
-      return emitError() << "lds_load_indices must only contain WaveSymbolAttr or WaveIndexSymbolAttr";
+      return emitError() << "lds_load_indices must only contain WaveSymbolAttr "
+                            "or WaveIndexSymbolAttr";
     }
 
     if (!llvm::all_of(lds_load_vector_sizes.getSymbols(),
                       llvm::IsaPred<WaveSymbolAttr, WaveIndexSymbolAttr>)) {
-      return emitError() << "lds_load_vector_sizes must only contain WaveSymbolAttr or WaveIndexSymbolAttr";
+      return emitError() << "lds_load_vector_sizes must only contain "
+                            "WaveSymbolAttr or WaveIndexSymbolAttr";
     }
 
     if (!llvm::all_of(global_store_indices.getSymbols(),
                       llvm::IsaPred<WaveSymbolAttr, WaveIndexSymbolAttr>)) {
-      return emitError() << "global_store_indices must only contain WaveSymbolAttr or WaveIndexSymbolAttr";
+      return emitError() << "global_store_indices must only contain "
+                            "WaveSymbolAttr or WaveIndexSymbolAttr";
     }
 
     // Validate that mappings have at least one dimension
     if (ldsBaseRank == 0) {
       return emitError() << "LDS block must have at least one dimension";
     }
 
-    // Note: We cannot validate that the ranks match the original global memory tensor rank here
-    // because this attribute verification doesn't have access to the WriteOp's memory operand.
-    // This validation should be performed in the WriteOp's verifier where both the attribute
-    // and the memory operand type are available.
+    // Note: We cannot validate that the ranks match the original global memory
+    // tensor rank here because this attribute verification doesn't have access
+    // to the WriteOp's memory operand. This validation should be performed in
+    // the WriteOp's verifier where both the attribute and the memory operand
+    // type are available.
     //
-    // Additionally, data coverage verification (ensuring that the collective workgroup access
-    // pattern covers exactly the same elements before and after LDS promotion) should be
-    // performed in the WriteOp verifier where access to the original index mapping is available.
+    // Additionally, data coverage verification (ensuring that the collective
+    // workgroup access pattern covers exactly the same elements before and
+    // after LDS promotion) should be performed in the WriteOp verifier where
+    // access to the original index mapping is available.
   }
 
   return success();