change waves_per_block to DenseI32ArrayAttr to allow to use struct(params) in assemblyFormat

tgymnich · tgymnich · commit 7323d2d05fa6 · 2026-03-10T12:00:46.000Z
Signed-off-by: Tim Gymnich &lt;tim@gymni.ch&gt;
diff --git a/water/include/water/Dialect/Wave/IR/WaveAttrs.td b/water/include/water/Dialect/Wave/IR/WaveAttrs.td
@@ -237,18 +237,12 @@ def HardwareConstraintAttr : AttrDef<WaveDialect, "HardwareConstraint"> {
     configuration rather than fundamental hardware constraints.
   }];
   let parameters = (ins "unsigned":$threads_per_wave,
-                        OptionalArrayRefParameter<"unsigned">:$waves_per_block,
+                        OptionalParameter<"::mlir::DenseI32ArrayAttr">:$waves_per_block,
                         OptionalParameter<"::wave::WaveMmaKindAttr">:$mma_type,
                         OptionalParameter<"::mlir::DictionaryAttr">:$vector_shapes,
                         DefaultValuedParameter<"unsigned", "128">:$max_bits_per_load);
 
-  let assemblyFormat = [{
-    `<` `threads_per_wave` `=` $threads_per_wave
-    (`,` `waves_per_block` `=` `[` $waves_per_block^ `]`)?
-    (`,` `mma_type` `=` $mma_type^)?
-    (`,` `vector_shapes` `=` $vector_shapes^)?
-    (`,` `max_bits_per_load` `=` $max_bits_per_load^)? `>`
-  }];
+  let assemblyFormat = "`<` struct(params) `>`";
 
   let genVerifyDecl = 1;
 }
diff --git a/water/include/water/Dialect/Wave/IR/WaveInterfaces.h b/water/include/water/Dialect/Wave/IR/WaveInterfaces.h
@@ -558,8 +558,8 @@ class IndexExprsAnalysisInit {
       symbolConstraints;
 
   // Waves-per-block extracted from the hardware constraint or computed from
-  // wave constraints. Always stored here, even if copied from an attribute.
-  llvm::SmallVector<unsigned, 3> wavesPerBlock;
+  // wave constraints.
+  llvm::SmallVector<int32_t, 3> wavesPerBlock;
 };
 
 // Lattice for propagating index expressions across wave dialect operations.
diff --git a/water/include/water/Dialect/Wave/IR/WaveUtils.h b/water/include/water/Dialect/Wave/IR/WaveUtils.h
@@ -62,7 +62,7 @@ llvm::LogicalResult computeWavesPerBlockFromConstraints(
     const llvm::SmallDenseMap<wave::WaveSymbolAttr, wave::WaveConstraintAttr>
         &waveConstraints,
     wave::WaveHyperparameterAttr hyperparams,
-    llvm::SmallVectorImpl<unsigned> &wavesPerBlock);
+    llvm::SmallVectorImpl<int32_t> &wavesPerBlock);
 
 /// Permute the shape according to the mapping.
 void permuteShape(llvm::ArrayRef<wave::WaveSymbolAttr> shape,
diff --git a/water/include/water/c/Dialects.h b/water/include/water/c/Dialects.h
@@ -475,7 +475,7 @@ mlirAttributeIsAHardwareConstraintAttr(MlirAttribute attr);
 /// Creates a new HardwareConstraintAttr
 MLIR_CAPI_EXPORTED MlirAttribute mlirHardwareConstraintAttrGet(
     MlirContext mlirCtx, unsigned threadsPerWave, size_t wavesPerBlockSize,
-    unsigned *wavesPerBlock, MlirAttribute mmaType, MlirAttribute vectorShapes,
+    int32_t *wavesPerBlock, MlirAttribute mmaType, MlirAttribute vectorShapes,
     unsigned maxBitsPerLoad);
 
 /// Returns the typeID of a HardwareConstraintAttr.
@@ -486,7 +486,7 @@ MLIR_CAPI_EXPORTED unsigned
 mlirHardwareConstraintAttrGetThreadsPerWave(MlirAttribute attr);
 MLIR_CAPI_EXPORTED intptr_t
 mlirHardwareConstraintAttrGetNumWavesPerBlock(MlirAttribute attr);
-MLIR_CAPI_EXPORTED unsigned
+MLIR_CAPI_EXPORTED int32_t
 mlirHardwareConstraintAttrGetWavesPerBlockElem(MlirAttribute attr, intptr_t i);
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirHardwareConstraintAttrGetMmaType(MlirAttribute attr);
diff --git a/water/lib/CAPI/Dialects.cpp b/water/lib/CAPI/Dialects.cpp
@@ -516,7 +516,7 @@ bool mlirAttributeIsAHardwareConstraintAttr(MlirAttribute attr) {
 
 MlirAttribute
 mlirHardwareConstraintAttrGet(MlirContext mlirCtx, unsigned threadsPerWave,
-                              size_t wavesPerBlockSize, unsigned *wavesPerBlock,
+                              size_t wavesPerBlockSize, int32_t *wavesPerBlock,
                               MlirAttribute mmaType, MlirAttribute vectorShapes,
                               unsigned maxBitsPerLoad) {
   MLIRContext *ctx = unwrap(mlirCtx);
@@ -525,9 +525,14 @@ mlirHardwareConstraintAttrGet(MlirContext mlirCtx, unsigned threadsPerWave,
   auto vectorShapesAttr =
       llvm::cast_if_present<DictionaryAttr>(unwrap(vectorShapes));
 
+  DenseI32ArrayAttr wavesPerBlockAttr;
+  if (wavesPerBlockSize > 0)
+    wavesPerBlockAttr = DenseI32ArrayAttr::get(
+        ctx, llvm::ArrayRef(wavesPerBlock, wavesPerBlockSize));
+
   return wrap(wave::HardwareConstraintAttr::get(
-      ctx, threadsPerWave, llvm::ArrayRef(wavesPerBlock, wavesPerBlockSize),
-      mmaTypeAttr, vectorShapesAttr, maxBitsPerLoad));
+      ctx, threadsPerWave, wavesPerBlockAttr, mmaTypeAttr, vectorShapesAttr,
+      maxBitsPerLoad));
 }
 
 MlirTypeID mlirWHardwareConstraintAttrGetTypeID() {
@@ -539,14 +544,15 @@ unsigned mlirHardwareConstraintAttrGetThreadsPerWave(MlirAttribute attr) {
       .getThreadsPerWave();
 }
 intptr_t mlirHardwareConstraintAttrGetNumWavesPerBlock(MlirAttribute attr) {
-  return llvm::cast<wave::HardwareConstraintAttr>(unwrap(attr))
-      .getWavesPerBlock()
-      .size();
+  DenseI32ArrayAttr wpb =
+      llvm::cast<wave::HardwareConstraintAttr>(unwrap(attr)).getWavesPerBlock();
+  return wpb ? wpb.size() : 0;
 }
-unsigned mlirHardwareConstraintAttrGetWavesPerBlockElem(MlirAttribute attr,
-                                                        intptr_t i) {
+int32_t mlirHardwareConstraintAttrGetWavesPerBlockElem(MlirAttribute attr,
+                                                       intptr_t i) {
   return llvm::cast<wave::HardwareConstraintAttr>(unwrap(attr))
-      .getWavesPerBlock()[i];
+      .getWavesPerBlock()
+      .asArrayRef()[i];
 }
 MlirAttribute mlirHardwareConstraintAttrGetMmaType(MlirAttribute attr) {
   return wrap(
diff --git a/water/lib/Dialect/Wave/IR/WaveAttrs.cpp b/water/lib/Dialect/Wave/IR/WaveAttrs.cpp
@@ -708,12 +708,11 @@ WaveExprListAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 
 LogicalResult HardwareConstraintAttr::verify(
     function_ref<InFlightDiagnostic()> emitError, unsigned threadsPerWave,
-    ArrayRef<unsigned> wavesPerBlock, WaveMmaKindAttr mmaType,
+    DenseI32ArrayAttr wavesPerBlock, WaveMmaKindAttr mmaType,
     DictionaryAttr vectorShapes, unsigned maxBitsPerLoad) {
 
-  if (!(wavesPerBlock.empty() || wavesPerBlock.size() == 3))
-    return emitError() << "waves_per_block (" << wavesPerBlock
-                       << ") should have 3 elements";
+  if (wavesPerBlock && wavesPerBlock.size() != 3)
+    return emitError() << "waves_per_block should have 3 elements";
 
   if (vectorShapes) {
     for (NamedAttribute attr : vectorShapes) {
diff --git a/water/lib/Dialect/Wave/IR/WaveDialect.cpp b/water/lib/Dialect/Wave/IR/WaveDialect.cpp
@@ -315,10 +315,10 @@ verifyConstraints(ArrayAttr constraints,
   // verify consistency between wave constraints and waves_per_block
   // * If both wave constraints and waves_per_block are present, the computed
   // number of waves per dimension should match the waves_per_block attribute.
-  if (hardwareConstraint && !hardwareConstraint.getWavesPerBlock().empty() &&
+  if (hardwareConstraint && hardwareConstraint.getWavesPerBlock() &&
       !waveConstraints.empty()) {
-    llvm::ArrayRef<unsigned> wavesPerBlock =
-        hardwareConstraint.getWavesPerBlock();
+    llvm::ArrayRef<int32_t> wavesPerBlock =
+        hardwareConstraint.getWavesPerBlock().asArrayRef();
     for (auto &&[symbol, waveConstraint] : waveConstraints) {
       wave::WorkgroupConstraintAttr wgConstraint = workgroupConstraints[symbol];
       unsigned wgDim =
diff --git a/water/lib/Dialect/Wave/IR/WaveInterfaces.cpp b/water/lib/Dialect/Wave/IR/WaveInterfaces.cpp
@@ -889,11 +889,11 @@ wave::IndexExprsAnalysisInit::create(Location loc, Attribute constraintsAttr,
   // If waves_per_block is explicitly provided, copy it to storage. Note that we
   // have verified they match the result of dividing block tiles with wave tiles
   // previously.
-  if (!initObject.hardwareConstraint.getWavesPerBlock().empty()) {
+  if (initObject.hardwareConstraint.getWavesPerBlock()) {
     assert(initObject.hardwareConstraint.getWavesPerBlock().size() == 3 &&
            "expected waves_per_block to have 3 elements");
-    llvm::ArrayRef<unsigned> explicitWpb =
-        initObject.hardwareConstraint.getWavesPerBlock();
+    llvm::ArrayRef<int32_t> explicitWpb =
+        initObject.hardwareConstraint.getWavesPerBlock().asArrayRef();
     initObject.wavesPerBlock.assign(explicitWpb);
     return initObject;
   }
diff --git a/water/lib/Dialect/Wave/IR/WaveOps.cpp b/water/lib/Dialect/Wave/IR/WaveOps.cpp
@@ -870,7 +870,7 @@ void MmaSingleIndexExprBuilder::populate(
 // dimension.
 static llvm::LogicalResult
 populateMmaIndexingExpr(wave::WaveMmaKind kind, bool isAccumulator,
-                        llvm::ArrayRef<unsigned> wavesPerWorkgroup,
+                        llvm::ArrayRef<int32_t> wavesPerWorkgroup,
                         int64_t threadsPerWave, wave::WaveSymbolAttr mSymbol,
                         wave::WaveSymbolAttr nSymbol,
                         wave::WaveSymbolAttr kSymbol,
diff --git a/water/lib/Dialect/Wave/IR/WaveUtils.cpp b/water/lib/Dialect/Wave/IR/WaveUtils.cpp
@@ -133,7 +133,7 @@ LogicalResult wave::computeWavesPerBlockFromConstraints(
     const llvm::SmallDenseMap<wave::WaveSymbolAttr, wave::WaveConstraintAttr>
         &waveConstraints,
     wave::WaveHyperparameterAttr hyperparams,
-    SmallVectorImpl<unsigned> &wavesPerBlock) {
+    SmallVectorImpl<int32_t> &wavesPerBlock) {
   // Default to 1 wave per block for each dimension, this may be recomputed
   // later if the corresponding constraints are provided.
   wavesPerBlock.assign(/*NumElts=*/3, /*Elt=*/1);
@@ -168,7 +168,7 @@ LogicalResult wave::computeWavesPerBlockFromConstraints(
     int64_t numWaves = workgroupSize / waveSize;
     unsigned wgDim =
         static_cast<unsigned>(wgConstraint.getWorkgroupDim().getValue());
-    wavesPerBlock[wgDim] = static_cast<unsigned>(numWaves);
+    wavesPerBlock[wgDim] = static_cast<int32_t>(numWaves);
   }
 
   return success();
diff --git a/water/lib/Dialect/Wave/Transforms/LoweringPatterns.cpp b/water/lib/Dialect/Wave/Transforms/LoweringPatterns.cpp
@@ -1327,9 +1327,12 @@ static void warnIfReductionScopeMismatch(Operation *op, bool isBlockReduction) {
         llvm::dyn_cast<wave::HardwareConstraintAttr>(constraint);
     if (!hwConstraint)
       continue;
-    ArrayRef<unsigned> wavesPerBlock = hwConstraint.getWavesPerBlock();
-    unsigned totalWaves = 1;
-    for (unsigned w : wavesPerBlock)
+    if (!hwConstraint.getWavesPerBlock())
+      continue;
+    llvm::ArrayRef<int32_t> wavesPerBlock =
+        hwConstraint.getWavesPerBlock().asArrayRef();
+    int64_t totalWaves = 1;
+    for (int32_t w : wavesPerBlock)
       totalWaves *= w;
     if (isBlockReduction && totalWaves == 1) {
       op->emitWarning()
diff --git a/water/python/WaterExtensionNanobind.cpp b/water/python/WaterExtensionNanobind.cpp
@@ -784,15 +784,15 @@ struct PyHardwareConstraintAttr
     c.def_static(
         "get",
         [](unsigned threadsPerWave,
-           const std::optional<std::vector<unsigned>> &wavesPerBlock,
+           const std::optional<std::vector<int32_t>> &wavesPerBlock,
            std::optional<MlirAttribute> mmaType,
            std::optional<MlirAttribute> vectorShapes, unsigned maxBitsPerLoad,
            mlir::python::MLIR_BINDINGS_PYTHON_DOMAIN::DefaultingPyMlirContext
                context) {
-          unsigned *wavesPerBlockPtr = nullptr;
+          int32_t *wavesPerBlockPtr = nullptr;
           size_t wavesPerBlockSize = 0;
           if (wavesPerBlock.has_value()) {
-            wavesPerBlockPtr = const_cast<unsigned *>(wavesPerBlock->data());
+            wavesPerBlockPtr = const_cast<int32_t *>(wavesPerBlock->data());
             wavesPerBlockSize = wavesPerBlock->size();
           }
 
@@ -812,7 +812,7 @@ struct PyHardwareConstraintAttr
       return mlirHardwareConstraintAttrGetThreadsPerWave(self);
     });
     c.def_prop_ro("waves_per_block", [](MlirAttribute self) {
-      std::vector<unsigned> out;
+      std::vector<int32_t> out;
       intptr_t n = mlirHardwareConstraintAttrGetNumWavesPerBlock(self);
       out.reserve(n);
       for (intptr_t i = 0; i < n; ++i)
