verify vector_shapes consistency with WorkgroupConstraint and TilingConstraint

Signed-off-by: Tim Gymnich <tim@gymni.ch>

Author: tgymnich

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

@@ -335,6 +335,8 @@ verifyConstraints(ArrayAttr constraints,
 
   // verify TilingConstraint
   // * The number of tiles should be greater than or equal to one.
+  llvm::SmallDenseMap<wave::WaveSymbolAttr, int64_t> resolvedTilingSizes(
+      tilingConstraints.size());
   for (auto &&[symbol, constraint] : tilingConstraints) {
     std::optional<llvm::SmallVector<int64_t>> evaluated =
         wave::evaluateMapWithHyperparams(constraint.getTileSize().getMap(),
@@ -351,6 +353,7 @@ verifyConstraints(ArrayAttr constraints,
            "failed to resolve dimesion symbol");
 
     int64_t resolvedTileSize = evaluated->front();
+    resolvedTilingSizes[symbol] = resolvedTileSize;
     int64_t resolvedDim = resolvedDims->front();
     int64_t numTiles = resolvedDim / resolvedTileSize;
     if (numTiles < 1) {
@@ -359,6 +362,44 @@ verifyConstraints(ArrayAttr constraints,
     }
   }
 
+  // Verify consistency between vector_shapes and constraint tile sizes.
+  // * When no mma_type is set, vector_shapes should match the resolved tile
+  //   sizes from the corresponding WorkgroupConstraint/TilingConstraint.
+  // * When mma_type is present, vector_shapes derives from the MMA instruction
+  //   geometry and is expected to differ from the constraint tile sizes.
+  if (hardwareConstraint && hardwareConstraint.getVectorShapes() &&
+      !hardwareConstraint.getMmaType()) {
+    DictionaryAttr vectorShapes = hardwareConstraint.getVectorShapes();
+    for (NamedAttribute entry : vectorShapes) {
+      llvm::StringRef dimName = entry.getName();
+      int64_t vectorShapeSize =
+          llvm::cast<IntegerAttr>(entry.getValue()).getInt();
+
+      for (auto &&[symbol, size] : resolvedWorkgroupSizes) {
+        if (symbol.getName() != dimName)
+          continue;
+        if (size != vectorShapeSize) {
+          return emitError()
+                 << "vector_shapes entry '" << dimName << "' ("
+                 << vectorShapeSize
+                 << ") does not match workgroup constraint tile size (" << size
+                 << ") for dimension: " << symbol;
+        }
+      }
+
+      for (auto &&[symbol, size] : resolvedTilingSizes) {
+        if (symbol.getName() != dimName)
+          continue;
+        if (size != vectorShapeSize) {
+          return emitError() << "vector_shapes entry '" << dimName << "' ("
+                             << vectorShapeSize
+                             << ") does not match tiling constraint tile size ("
+                             << size << ") for dimension: " << symbol;
+        }
+      }
+    }
+  }
+
   return llvm::success();
 }
 

water/test/Dialect/Wave/attr-constraint-invalid.mlir

@@ -1,6 +1,6 @@
 // RUN: water-opt %s --allow-unregistered-dialect --water-test-wave-dialect-functions --split-input-file --verify-diagnostics
 
-// expected-error @below {{waves_per_block (1) should have 3 elements}}
+// expected-error @below {{waves_per_block should have 3 elements}}
 #hw_constraint = #wave.hardware_constraint<threads_per_wave = 64,
                                            waves_per_block = [1],
                                            mma_type = #wave.mma_kind<f32_16x16x16_f16>,
@@ -252,3 +252,19 @@ func.func private @test_waves_per_block_mismatch_multi_dim() attributes { wave.h
 func.func private @test_waves_per_block_mismatch_y_dim() attributes { wave.hyperparameters = #hyperparams_wpb3, wave.constraints = [#wg_constraint_wpb3, #wv_constraint_wpb3, #hw_constraint_wpb3] }
 
 // -----
+
+#hyperparams_vs1 = #wave.hyperparameters<{M = 64, BLOCK_M = 32}>
+#wg_constraint_vs1 = #wave.workgroup_constraint<dim = <"M">, tile_size = <[#wave.symbol<"BLOCK_M">] -> (BLOCK_M)>, workgroup_dim = <x>>
+#hw_constraint_vs1 = #wave.hardware_constraint<threads_per_wave = 64, vector_shapes = {M = 16}>
+// expected-error @below {{vector_shapes entry 'M' (16) does not match workgroup constraint tile size (32) for dimension: #wave.symbol<"M">}}
+func.func private @test_vector_shapes_wg_mismatch() attributes { wave.hyperparameters = #hyperparams_vs1, wave.constraints = [#wg_constraint_vs1, #hw_constraint_vs1] }
+
+// -----
+
+#hyperparams_vs2 = #wave.hyperparameters<{K = 1024, BLOCK_K = 128}>
+#tl_constraint_vs2 = #wave.tiling_constraint<dim = <"K">, tile_size = <[#wave.symbol<"BLOCK_K">] -> (BLOCK_K)>>
+#hw_constraint_vs2 = #wave.hardware_constraint<threads_per_wave = 64, vector_shapes = {K = 32}>
+// expected-error @below {{vector_shapes entry 'K' (32) does not match tiling constraint tile size (128) for dimension: #wave.symbol<"K">}}
+func.func private @test_vector_shapes_tiling_mismatch() attributes { wave.hyperparameters = #hyperparams_vs2, wave.constraints = [#tl_constraint_vs2, #hw_constraint_vs2] }
+
+// -----