[Encoding] Add `iree_encoding.dim` op and reification #23311

jtuyls · 2026-01-28T17:39:56Z

Implements phase 3 for adding specialization support on dynamic values for data-tiling: #22370. This adds the iree_encoding.dim operation and reification patterns to query encoding dimensions. For example:

// Encoding with dynamic M dimension (iteration_sizes = [?, 64, 128])
// encoding_dims has one element: the M value
%encoded = iree_encoding.set_encoding %input encoding_dims{%m}
    : tensor<?x128xf32> -> tensor<?x128xf32, #encoding>

// Query the first (and only) dynamic encoding dimension
%m_value = iree_encoding.dim %encoded[0] : tensor<?x128xf32, #encoding>

Assisted-by: Claude

Signed-off-by: Jorn Tuyls <[email protected]>

hanhanW

I feel that adding the interface may be over-design. What you're trying to do in this PR can just have a single canonicalization pattern? Below is an implementation example from Claude.

I think the main question is: why do we need this interface? In practice, the encodings are either come from SetEncoding op or frontend. Is it because you need such support for Flow::EncodeOp and Stream::EncodeOp as well?

  struct ReifyEncodingDim : public OpRewritePattern<DimOp> {
    LogicalResult matchAndRewrite(DimOp dimOp,
                                  PatternRewriter &rewriter) const override {
      auto result = dyn_cast<OpResult>(dimOp.getSource());
      if (!result)
        return failure();

      Operation *producer = result.getOwner();
      int64_t dimIndex = dimOp.getConstantIndex();

      // Source: set_encoding directly provides encoding dims.
      if (auto setEnc = dyn_cast<SetEncodingOp>(producer)) {
        ValueRange encodingDims = setEnc.getEncodingDims();
        if (dimIndex >= encodingDims.size())
          return failure();
        rewriter.replaceOp(dimOp, encodingDims[dimIndex]);
        return success();
      }

      // Pass-through: tensor.cast forwards to source.
      if (auto castOp = dyn_cast<tensor::CastOp>(producer)) {
        rewriter.replaceOpWithNewOp<DimOp>(dimOp, castOp.getSource(),
                                           dimIndex);
        return success();
      }

      // Pass-through: DPS ops forward to tied init.
      if (auto dpsOp = dyn_cast<DestinationStyleOpInterface>(producer)) {
        if (auto *tiedInit = dpsOp.getTiedOpOperand(result)) {
          rewriter.replaceOpWithNewOp<DimOp>(dimOp, tiedInit->get(),
                                             dimIndex);
          return success();
        }
      }

      return failure();
    }
  };

btw, please expand the context in PR description. I was not aware of an interface is added until I review the code. Here is a good guidance: https://google.github.io/eng-practices/review/developer/cl-descriptions.html#informative

hanhanW · 2026-01-29T22:23:01Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td

+    This interface enables reification of `iree_encoding.encoding_dim` operations
+    by tracing through producer chains to find where encoding dimension values
+    were originally captured.


Suggested change

This interface enables reification of `iree_encoding.encoding_dim` operations

by tracing through producer chains to find where encoding dimension values

were originally captured.

This interface enables reification of `iree_encoding.dim` operations by tracing

through producer chains to find where encoding dimension values were

originally captured.

hanhanW · 2026-01-29T22:38:54Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingOps.td

+    - `set_encoding` implements `EncodingDimReificationInterface` and returns
+      the corresponding `encoding_dims` value
+    - `tensor.cast` and DPS ops (like `linalg.fill`, `linalg.generic`) forward
+      the query to their source/init operands


Suggested change

- `set_encoding` implements `EncodingDimReificationInterface` and returns

the corresponding `encoding_dims` value

- `tensor.cast` and DPS ops (like `linalg.fill`, `linalg.generic`) forward

the query to their source/init operands

- `set_encoding` implements `EncodingDimReificationInterface` and returns

the corresponding `encoding_dims` value.

- `tensor.cast` and DPS ops (like `linalg.fill`, `linalg.generic`) forward

the query to their source/init operands.

hanhanW · 2026-01-29T22:46:29Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingOps.td

+  let results = (outs Index:$result);
+
+  let assemblyFormat = [{
+    attr-dict $source `[` $index `]` `:` type($source)


Making attribute list after source and index seems more common?

Suggested change

attr-dict $source `[` $index `]` `:` type($source)

$source `[` $index `]` attr-dict `:` type($source)

hanhanW · 2026-01-29T22:54:33Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td

+        /*methodName=*/"reifyEncodingDim",
+        /*args=*/(ins
+          "::mlir::OpBuilder &":$builder,
+          "unsigned":$resultIndex,


Do we pass OpResult instead? It usually provides more information and passing it is not expensive as my understanding is that it is a pointer-like value.

I don't have a full picture about how you'd use the interface, so I'll leave the decision to you.

hanhanW · 2026-01-29T22:56:54Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td

+          - Success with the value if the dimension can be resolved directly
+          - Failure if the operation cannot directly provide the value
+            (caller should use `getEncodingDimSource` to trace through)


Suggested change

- Success with the value if the dimension can be resolved directly

- Failure if the operation cannot directly provide the value

(caller should use `getEncodingDimSource` to trace through)

- Success with the value if the dimension can be resolved directly.

- Failure if the operation cannot directly provide the value.

(caller should use `getEncodingDimSource` to trace through)

Do we need FailureOr? Do we just follow the other method that returns either a Value or null?

caller should use getEncodingDimSource to trace through

Can you collaborate a bit more? Does it mean that caller uses wrong method if it returns failure?

hanhanW · 2026-01-29T23:03:19Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingPatterns.h

It is weird to see *Patterns under IR/. Are they only used by canonicalization patterns? If so, can you move them to EncodingOps.cpp?

hanhanW · 2026-01-29T23:03:39Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingPatterns.cpp

+/// 2. Operations that forward encoding dims from a source (like tensor.cast):
+///    The pattern calls getEncodingDimSource() and creates a new dim op on
+///    that source.
+///


I'd drop this blank comment.

hanhanW · 2026-01-29T23:06:39Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingPatterns.cpp

+    OpResult result = dyn_cast<OpResult>(dimOp.getSource());
+    if (!result) {
+      return failure();
+    }


Suggested change

OpResult result = dyn_cast<OpResult>(dimOp.getSource());

if (!result) {

return failure();

}

auto result = dyn_cast<OpResult>(dimOp.getSource());

if (!result) {

return failure();

}

Please also replace return failure() with more meaningful message if possible. I.e., return rewriter.notifyMatchFailure(...). The error message is also a self-comment, which looks better to me.

hanhanW · 2026-01-29T23:13:34Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingPatterns.cpp

+    // Verify encodings match.
+    auto resultType = dyn_cast<RankedTensorType>(result.getType());
+    auto initType = dyn_cast<RankedTensorType>(tiedInit->get().getType());
+    if (!resultType || !initType) {
+      return failure();
+    }
+
+    if (resultType.getEncoding() != initType.getEncoding()) {
+      return failure();
+    }


My intuition told me that it should be checked by the interface; yes, I confirmed it: https://github.com/llvm/llvm-project/blob/52dfcab327fe959074563603b6ebaaed314e9677/mlir/lib/Interfaces/DestinationStyleOpInterface.cpp#L51-L59

for (OpOperand *opOperand : outputTensorOperands) { OpResult result = dstStyleOp.getTiedOpResult(opOperand); if (result.getType() != opOperand->get().getType()) return op->emitOpError("expected type of operand #") << opOperand->getOperandNumber() << " (" << opOperand->get().getType() << ")" << " to match type of corresponding result (" << result.getType() << ")"; }

They have the same type, which indicates that they have the same encoding.

hanhanW · 2026-01-29T23:17:58Z

compiler/src/iree/compiler/Dialect/Encoding/IR/EncodingInterfaces.td

Why don't we have a unified interface method? Then the next question is, why do we need the interface? Can't it just be a single method?

(I may be missing how it is used in other places. I haven't reached to that state yet.)

jtuyls · 2026-01-30T09:33:13Z

I feel that adding the interface may be over-design. What you're trying to do in this PR can just have a single canonicalization pattern?

Imo, the interface is appropriate as for the cost of a bit additional code we get:

Separation of concerns: No op-specific special casing in the transformation pattern. New op support can be added through the the interface instead of new if blocks for each supported op (SetEncodingOp, tensor.cast, DPS ops, and possibly expand_shape, collapse_shape, extract_slice, etc.). So, we separate the transformation pattern from the op specific logic.
External models / downstream extensibility: Projects building on IREE can extend to their custom ops by registering an external model.
Explicit contract: The interface definition documents clearly what ops must provide.

Per my understanding, interfaces exist precisely to avoid encoding op-specific knowledge into transformation logic, and this keeps future additions isolated.

jtuyls requested review from Max191, benvanik, bjacob and hanhanW as code owners January 28, 2026 17:39

[Encoding] Add dim op and reification

26c689c

Signed-off-by: Jorn Tuyls <[email protected]>

jtuyls force-pushed the users/jtuyls/phase-3-encoding-dim-op branch from 38700b3 to 26c689c Compare January 29, 2026 14:30

hanhanW requested changes Jan 29, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Encoding] Add `iree_encoding.dim` op and reification #23311

[Encoding] Add `iree_encoding.dim` op and reification #23311

Uh oh!

jtuyls commented Jan 28, 2026

Uh oh!

hanhanW left a comment

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

hanhanW Jan 29, 2026

Uh oh!

jtuyls commented Jan 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

	attr-dict $source `[` $index `]` `:` type($source)
	$source `[` $index `]` attr-dict `:` type($source)

[Encoding] Add iree_encoding.dim op and reification #23311

Are you sure you want to change the base?

[Encoding] Add iree_encoding.dim op and reification #23311

Uh oh!

Conversation

jtuyls commented Jan 28, 2026

Uh oh!

hanhanW 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!

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!

jtuyls commented Jan 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[Encoding] Add `iree_encoding.dim` op and reification #23311

[Encoding] Add `iree_encoding.dim` op and reification #23311