[water] Add support for vector types in wave.iterate and wave.yield operations #625

tyb0807 · 2025-12-23T01:57:43Z

Stacked PRs, do not merge.

The operations now accept both WaveTensorInRegister
(before conversion) and VectorOfAnyType (after conversion).

Updated type compatibility and verification logic to handle both tensor
and vector type combinations appropriately.

Fixes #624.

Implements elements per thread propagation for MMA operations. Fixes iree-org#608. Signed-off-by: tyb0807 <sontuan.vu@amd.com>

Changes: - ReadOp: Only propagate attribute to result (register), ignore memory - WriteOp: Only validate/propagate with register operand, ignore memory This fixes false positives where memory resharding was incorrectly flagged as propagation errors. Fixes iree-org#622. Signed-off-by: tyb0807 <sontuan.vu@amd.com>

The operations now accept both WaveTensorInRegister (before conversion) and VectorOfAnyType (after conversion). Updated type compatibility and verification logic to handle both tensor and vector type combinations appropriately. Fixes iree-org#624. Signed-off-by: tyb0807 <sontuan.vu@amd.com>

ftynse · 2025-12-23T08:25:27Z

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

-    Arg<Variadic<WaveTensorType>, "Carried values">:$iter_args,
-    Arg<Variadic<WaveTensorType>, "Captured values">:$captures
+    // Accept both WaveTensorType (before PropagateElementsPerThread) and AnyVectorOfAnyRank (after)
+    Arg<Variadic<AnyTypeOf<[WaveTensorType, AnyVectorOfAnyRank]>>, "Carried values">:$iter_args,


Why can't this be WaveTensorInRegisters? That constraint already accepts the the tensors with no address space, tensors in register address space and 1D vectors. And we most likely don't want any vector of any rank here, which would include scalable, 0d and other nonsense.

This is because WaveTensorInRegisters doesn't work with Variadic. I think because it's a TypeConstraint (or something like that) and not a Type.

Add a comment explaining that, this is something that should be fixed upstream eventually

And we still don't want any vector of any rank here. We specifically want a 1D vector. It is also significantly easier to maintain if you somehow created a named tabelgen entity for it.

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

tgymnich · 2025-12-23T10:34:50Z

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

      std::optional<int64_t> value = hyper.getSymbolValue(name);
 #ifndef NDEBUG
      if (!value) {
-        llvm::errs() << "symbol: " << name << "\n";


why remove this?

Indeed, this is extra output that will be printed before the assertion, not a debug output.

ftynse

@tgymnich , when reviewing stacked PRs (the target branch is not main), click on the last commit and only review that one to avoid making comments on things that should be addressed in other PRs.

tgymnich · 2025-12-23T16:41:41Z

@tgymnich , when reviewing stacked PRs (the target branch is not main), click on the last commit and only review that one to avoid making comments on things that should be addressed in other PRs.

@ftynse could we instead just change the base to the PR below in the stack?

ftynse · 2025-12-24T09:08:05Z

I can adapt to the style folks use. It did accidentally click on "squash and merge" in a thus stacked PR before, polluted the other branch and had to do a bunch of force-pushing and PR re-opening to fix that.

ftynse · 2025-12-29T08:32:15Z

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

-    Arg<Variadic<WaveTensorType>, "Carried values">:$iter_args,
-    Arg<Variadic<WaveTensorType>, "Captured values">:$captures
+    // Accept both WaveTensorType (before PropagateElementsPerThread) and AnyVectorOfAnyRank (after)
+    Arg<Variadic<AnyTypeOf<[WaveTensorType, AnyVectorOfAnyRank]>>, "Carried values">:$iter_args,


And we still don't want any vector of any rank here. We specifically want a 1D vector. It is also significantly easier to maintain if you somehow created a named tabelgen entity for it.

ftynse · 2025-12-29T08:33:35Z

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

-                                       llvm::cast<wave::WaveTensorType>(rhs),
-                                       /*includeAddressSpace=*/true)
-      .succeeded();
+  // Handle both WaveTensorType and VectorType combinations


I usually don't comment on this, but please just take the habit (or configure your code generation assistant) to use full stop at the end of a sentence in comments.

ftynse · 2025-12-29T08:35:15Z

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

-            "result #" + istr, resultTensor, allDims)))
-      return mlir::failure();
+
+    // Both are wave tensors - use existing shape verification logic


"existing" doesn't make sense in the standalone comment when reading the code. It only makes sense when reading the diff.

ftynse · 2025-12-29T08:36:28Z

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

      std::optional<int64_t> value = hyper.getSymbolValue(name);
 #ifndef NDEBUG
      if (!value) {
-        llvm::errs() << "symbol: " << name << "\n";


Indeed, this is extra output that will be printed before the assertion, not a debug output.

ftynse · 2025-12-29T08:37:21Z

water/test/Dialect/Wave/ops.mlir

+  return
+}
+
+// CHECK-LABEL: @iterate_multidim_vectors


We don't want these.

ftynse · 2025-12-29T08:37:53Z

water/test/Dialect/Wave/ops.mlir

+// CHECK-LABEL: @iterate_vector_captures
+func.func @iterate_vector_captures() {
+  %iter_arg = arith.constant dense<1.0> : vector<8xf32>
+  %capture = arith.constant dense<2.0> : vector<4xf16>
+
+  // CHECK: wave.iterate @I iter_args(%{{.*}}) captures(%{{.*}})
+  %result = wave.iterate @I iter_args(%iter_arg) captures(%capture) {
+  ^bb0(%in_arg: vector<8xf32>, %cap: vector<4xf16>):
+    // CHECK: wave.yield %{{.*}} : vector<8xf32>
+    wave.yield %in_arg : vector<8xf32>
+  } : (vector<8xf32>, vector<4xf16>) -> (vector<8xf32>)
+  return
+}


Why do we need this and what does it test?

ftynse · 2025-12-29T08:38:29Z

water/test/Dialect/Wave/propagate-elements-per-thread.mlir

+// CHECK-LABEL: @iterate_with_vectors_after_ept
+func.func @iterate_with_vectors_after_ept(%mem: !wave.tensor<[@M] of f32, <global>>)
+  attributes {wave.hyperparameters = #wave.hyperparameters<{M = 128, I = 4}>,
+              wave.constraints = [#wave.hardware_constraint<threads_per_wave = 64, waves_per_block = [1, 1, 1], mma_type = #wave.mma_kind<f32_32x32x8_f16>, vector_shapes = {M = 1, N = 1, K = 8}, max_bits_per_load = 128>]} {


Do we need all of these constraints?

ftynse · 2025-12-29T08:39:01Z

water/test/Dialect/Wave/propagate-elements-per-thread.mlir

+  // CHECK: wave.iterate @I iter_args({{.*}})
+  %result = wave.iterate @I iter_args(%init) {
+  ^bb0(%arg: !wave.tensor<[@M] of f32, <register>>):
+    // Simple operation within the loop - should also work with vectors


At which point maybe this should // CHECK that it actually does?

ftynse · 2025-12-29T08:39:11Z

water/test/Dialect/Wave/propagate-elements-per-thread.mlir

+    wave.yield %doubled : !wave.tensor<[@M] of f32, <register>>
+  } : (!wave.tensor<[@M] of f32, <register>>) -> (!wave.tensor<[@M] of f32, <register>>)
+
+  // Write should also work with the vector result


How do we know it does?

tyb0807 added 3 commits December 23, 2025 01:58

Add WaveElementsPerThreadOpInterface to MmaOp

01abfb0

Implements elements per thread propagation for MMA operations. Fixes iree-org#608. Signed-off-by: tyb0807 <sontuan.vu@amd.com>

tyb0807 requested a review from ftynse December 23, 2025 01:57

ftynse reviewed Dec 23, 2025

View reviewed changes

tgymnich reviewed Dec 23, 2025

View reviewed changes

ftynse reviewed Dec 23, 2025

View reviewed changes

ftynse reviewed Dec 29, 2025

View reviewed changes

[water] Add support for vector types in wave.iterate and wave.yield operations #625

Are you sure you want to change the base?

[water] Add support for vector types in wave.iterate and wave.yield operations #625

Uh oh!

Conversation

tyb0807 commented Dec 23, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ftynse left a comment

Choose a reason for hiding this comment

Uh oh!

tgymnich commented Dec 23, 2025

Uh oh!

ftynse commented Dec 24, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants