[water] add wave.permute op #759

tgymnich · 2026-01-20T12:33:42Z

Add wave.permute op that permutes dimensions of a WaveTensorInRegister.
Lowers to a no-op, but modifies propagation of index expressions by permuting the strides according to target_shape.

martin-luecke

wave.permute uses the new CompatibleOperandsAndResultsIgnoreShapeOpTrait as verifier for its input and return values, which ignores the shape entirely.
But for a permute the rank should match, and the target_shape should be a permutation of the input symbols. I think we should enforce this in the verifier, otherwise we can end up with mismatched ranks or symbol sets that later assert (e.g., zip_equal) or mis‑infer.

Also, I think I found a couple of propagation issues I’ve flagged inline.

martin-luecke · 2026-01-23T09:02:39Z

wave_lang/kernel/wave/mlir_converter/water_emitter.py

    "iterate": IterateOp,
    "output": YieldOp,
    "write": WriteOp,
+    "permute": PermuteOp,


When adding it here please add a permute op to one of the pywave → wave dialect tests or create a new test for it so we know it actually works

added tests

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

water/test/Dialect/Wave/ops-invalid.mlir

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

martin-luecke · 2026-01-23T10:52:34Z

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

+  resultType = WaveTensorType::get(
+      getContext(), targetShape, /*fully_specified=*/true,
+      resultType ? resultType.getElementType() : inputType.getElementType(),
+      resultType ? resultType.getAddressSpace() : inputType.getAddressSpace());


I think this ternary can lead to loss of information when resultType.getAddressSpace() == wave::WaveAddressSpace::Unspecified and inputType actually has a specified address space

Also propagateForward only checks that resultType matches target_shape (when the result is already fully specified). It never validates that target_shape is a permutation of the input shape. That means you can end up setting a fully‑specified result whose rank/symbols don’t correspond to the input, and the index‑expr code later assumes equal sizes. Let's add an explicit input<>target permutation check (or enforce it in the op verifier)?

This definitely should be in the op verifier, but we may also want to have another diagnostic here. It may be better for UX to say "inference resulted in shape conflict" than just say "shape conflict" on shape-less input DSL.

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

water/include/water/Dialect/Wave/IR/WaveInterfaces.h

ftynse · 2026-01-24T19:10:56Z

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

+    compilation. At lowering time, the operation is a pass-through since the
+    actual data layout in registers remains unchanged - only the interpretation
+    of which dimension each element belongs to changes.


I have always found this concept of pass-through permutation difficult to grasp, can this be illustrated with an example?

ftynse · 2026-01-24T19:11:46Z

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

+  }];
+  let arguments = !con((ins
+    Arg<WaveTensorInRegister, "Value to permute">:$value,
+    Arg<WaveSymbolArrayAttr , "Target dimension ordering">:$target_shape


Do we need this attribute? When the result type is a tensor, it duplicates its shape. Is there a situation where we need it when the result type was lowered to a vector?

Good idea. Less to verify. We shouldn't need the shape downstream.

ftynse · 2026-01-24T19:14:07Z

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

+  // If result is already fully specified, verify it matches target_shape.
+  if (resultType && resultType.getFullySpecified()) {
+    ArrayRef<WaveSymbolAttr> resultShape = resultType.getShape();
+    if (resultShape.size() != targetShape.size()) {
+      errs << "result shape rank (" << resultShape.size()
+           << ") does not match target_shape rank (" << targetShape.size()
+           << ")";
+      return failure();
+    }
+    for (auto [i, expected, actual] :
+         llvm::enumerate(targetShape, resultShape)) {
+      if (expected != actual) {
+        errs << "result shape dimension #" << i << " (" << actual
+             << ") does not match target_shape (" << expected << ")";
+        return failure();
+      }
+    }


We have a helper function for this, checkPropagateShapeConflict.

ftynse · 2026-01-24T19:14:41Z

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

+  resultType = WaveTensorType::get(
+      getContext(), targetShape, /*fully_specified=*/true,
+      resultType ? resultType.getElementType() : inputType.getElementType(),
+      resultType ? resultType.getAddressSpace() : inputType.getAddressSpace());


This definitely should be in the op verifier, but we may also want to have another diagnostic here. It may be better for UX to say "inference resulted in shape conflict" than just say "shape conflict" on shape-less input DSL.

ftynse · 2026-01-24T19:15:21Z

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

+      resultType ? resultType.getElementType() : inputType.getElementType(),
+      resultType ? resultType.getAddressSpace() : inputType.getAddressSpace());
+
+  return ChangeResult::Change;


We need to check if the result type actually changes. Otherwise the analysis may never converge since we always indicate change. Does it?

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

ftynse · 2026-01-24T19:17:05Z

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

+  // Verify result shape matches target_shape.
+  if (resultShape.size() != targetShape.size()) {
+    errs << "result shape rank (" << resultShape.size()
+         << ") does not match target_shape rank (" << targetShape.size() << ")";
+    return failure();
+  }
+  for (auto [i, expected, actual] : llvm::enumerate(targetShape, resultShape)) {
+    if (expected != actual) {
+      errs << "result shape dimension #" << i << " (" << actual
+           << ") does not match target_shape (" << expected << ")";
+      return failure();
+    }
+  }


Same as above, we should have a helper doing this.

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

Signed-off-by: Tim Gymnich <[email protected]>

martin-luecke

With the changes this LGTM with a few nits still to address

martin-luecke · 2026-01-30T08:41:57Z

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

+  resultShapeSet.insert_range(resultType.getShape());
+
+  for (auto inputDim : inputType.getShape()) {
+    auto [_, inserted] = resultShapeSet.insert(inputDim);


Why do you not use contains() to check if the inputDim is in the set? Seems more idiomatic
e.g.
if (!resultShapeSet.contains(inputDim)) {

martin-luecke · 2026-01-30T08:48:35Z

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

+  if (inputType.getFullySpecified()) {
+    std::string errorMessage;
+    llvm::raw_string_ostream errs(errorMessage);
+    if (failed(validatePermutationInput(inputType, resultType, errs))) {
+      return emitOpError() << errorMessage;
+    }
+  }


If you pass in the op into validatePermutationInput as well, you wouldn't need the string here and could emit the error directly from inside it.
I think the MLIR idiomatic approach to passing error messages is to construct InFlightDiagnostic from the op and then compose it within validatePermutationInput.
However, this would require calling abandon() on it in case there was no error. I think the first approach is cleaner

This is due to how propagateForward and propagateBackward emit errors. If we directly emit here it will mess up the order of the messages. Might be something to fix for later, e.g. by passing a InFlightDiagnostic and adding notes.

martin-luecke · 2026-01-30T09:01:28Z

water/test/Dialect/Wave/infer-types.mlir

+  // CHECK: !wave.tensor<[@B, @M, @N] of f32, <register>> to !wave.tensor<[@M, @N, @B] of f32, <register>>
+  wave.permute %a : !wave.tensor<[@B, @M, @N] of f32, <register>> to !wave.tensor<[@M, @N, @B] of f32, <register>>
+  return
+}
+
+// CHECK-LABEL: @propagate_permute_2d
+func.func @propagate_permute_2d(%a: !wave.tensor<[@M, @N] of f16, <register>>) {
+  // CHECK: !wave.tensor<[@M, @N] of f16, <register>> to !wave.tensor<[@N, @M] of f16, <register>>
+  wave.permute %a : !wave.tensor<[@M, @N] of f16, <register>> to !wave.tensor<[@N, @M] of f16, <register>>
+  return


These CHECK lines only check parsing of the op, I don't see anything regarding type inference.
Could these be just removed in favor of the tests in ops.mlir or do they test anything in addition?

I think we can test the following:

%0 = wave.negate %arg0 : @A, @B, @C -> any wave.permute %0 any to @M, @N, @K

this will not fail the verifier, but the type inference should discover and report the conflict.

yes, this does not really propagate anything anymore, but its good to keep to make sure it does not fail. I'll also add the negative case.

ftynse

Nice, LGTM % nits

ftynse · 2026-01-30T09:31:18Z

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

+// PermuteOp
+//-----------------------------------------------------------------------------
+
+static LogicalResult validatePermutationInput(WaveTensorType inputType,


Nit please document functions

ftynse · 2026-01-30T09:32:28Z

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

+  // If result / input is a vector (post-lowering phase), skip wave tensor
+  // checks.


Nit: we want to verify element types for vectors as well. I added a helper function recently.

ftynse · 2026-01-30T09:35:00Z

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

+  }
+
+  // Result type is already specified, propagate it.
+  return detail::propagateShapeInformation(resultType, resultType,


Does this do anything useful? It looks like this will always succeed and the resultType should already be initialized.

ftynse · 2026-01-30T09:38:15Z

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

+  if (srcShape.size() != targetShape.size()) {
+    emitError() << "source shape rank (" << srcShape.size()
+                << ") does not match target shape rank (" << targetShape.size()
+                << ")";
+    return IndexExprsLatticeStorage::top();
+  }


I think this should be an assertion. We check in the verifier that shapes have equal rank when present, and we also have a normal form precondition for this entire analysis that types are fully specified. Try it if needed, it shouldn't be possible to produce this error message.

ftynse · 2026-01-30T09:40:04Z

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

+    // If the target or source mapping is not found, we cannot propagate the
+    // index expression.
+    if (srcMappingIt == symbolToMapping.end()) {


Can this happen in IR that passed the verifier and satisfied the full types normal form? If not, this should be an assertion. Same below.

If it can, maybe we should support some sort of partial propagation only for symbols that are present in hopes that other symbols show up later as we keep propagating.

ftynse · 2026-01-30T09:41:31Z

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

+  IndexExprsLatticeStorage permuted = permuteIndexExprsStrides(
+      operandExprs[0], srcShape, targetShape, getContext(), emitError);
+
+  permuted = permuted.keepOnlySymbols(resultType.getShape());


Is this needed? AFAIU, we should have strictly the same symbols before and after.

you are right. this is already handled.

ftynse · 2026-01-30T09:42:42Z

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

+  IndexExprsLatticeStorage permuted = permuteIndexExprsStrides(
+      resultExprs[0], resultShape, srcShape, getContext(), emitError);
+
+  permuted = permuted.keepOnlySymbols(srcShape);


Same as above

ftynse · 2026-01-30T09:46:39Z

water/test/Dialect/Wave/infer-types.mlir

+  // CHECK: !wave.tensor<[@B, @M, @N] of f32, <register>> to !wave.tensor<[@M, @N, @B] of f32, <register>>
+  wave.permute %a : !wave.tensor<[@B, @M, @N] of f32, <register>> to !wave.tensor<[@M, @N, @B] of f32, <register>>
+  return
+}
+
+// CHECK-LABEL: @propagate_permute_2d
+func.func @propagate_permute_2d(%a: !wave.tensor<[@M, @N] of f16, <register>>) {
+  // CHECK: !wave.tensor<[@M, @N] of f16, <register>> to !wave.tensor<[@N, @M] of f16, <register>>
+  wave.permute %a : !wave.tensor<[@M, @N] of f16, <register>> to !wave.tensor<[@N, @M] of f16, <register>>
+  return


I think we can test the following:

%0 = wave.negate %arg0 : @A, @B, @C -> any wave.permute %0 any to @M, @N, @K

this will not fail the verifier, but the type inference should discover and report the conflict.

Signed-off-by: Tim Gymnich <[email protected]>

tgymnich linked an issue Jan 20, 2026 that may be closed by this pull request

[water] Implement wave.permute #724

Open

tgymnich force-pushed the tim/permute-op branch 4 times, most recently from 474d06e to 9dd20e3 Compare January 22, 2026 14:24

tgymnich marked this pull request as ready for review January 22, 2026 14:27

tgymnich force-pushed the tim/permute-op branch from 9dd20e3 to dda879e Compare January 22, 2026 14:30

tgymnich requested review from martin-luecke and tyb0807 January 22, 2026 14:35

tgymnich force-pushed the tim/permute-op branch from dda879e to bb65e50 Compare January 22, 2026 15:56

martin-luecke requested changes Jan 23, 2026

View reviewed changes

ftynse reviewed Jan 24, 2026

View reviewed changes

ftynse reviewed Jan 26, 2026

View reviewed changes

water/lib/Dialect/Wave/IR/WaveOps.cpp Outdated Show resolved Hide resolved

water/lib/Dialect/Wave/IR/WaveOps.cpp Show resolved Hide resolved

water/lib/Dialect/Wave/IR/WaveOps.cpp Outdated Show resolved Hide resolved

tgymnich force-pushed the tim/permute-op branch 4 times, most recently from f0bcf95 to 47200f9 Compare January 28, 2026 09:54

[water] add wave.permute op

b96f4d3

Signed-off-by: Tim Gymnich <[email protected]>

tgymnich force-pushed the tim/permute-op branch from 47200f9 to b96f4d3 Compare January 28, 2026 10:02

update test

db6a5f3

Signed-off-by: Tim Gymnich <[email protected]>

tgymnich requested review from ftynse and martin-luecke January 29, 2026 12:37

martin-luecke approved these changes Jan 30, 2026

View reviewed changes

ftynse approved these changes Jan 30, 2026

View reviewed changes

tgymnich added 2 commits January 30, 2026 13:18

code review

3b5591a

Signed-off-by: Tim Gymnich <[email protected]>

fix mlir_exporter

3116fcf

Signed-off-by: Tim Gymnich <[email protected]>

		// If result / input is a vector (post-lowering phase), skip wave tensor
		// checks.

[water] add wave.permute op #759

Are you sure you want to change the base?

[water] add wave.permute op #759

Uh oh!

Conversation

tgymnich commented Jan 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

martin-luecke left a comment

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!

Uh oh!

Uh oh!

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!

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

martin-luecke left a comment

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!

ftynse left a comment

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!

tgymnich commented Jan 20, 2026 •

edited

Loading