Region-based shape inference: how it interacts with standard MLIR passes (repros + open questions)

[fhanuman]

Summary

In the hipsr direction, each op is designed to compute its own output shape inside a small region attached to the op — a shape_region (filled in by populateShapeRegion) ending in hipsr.shape_yield. That region is meant to be "sealed": it may read the op's own operands, but nothing else in the surrounding function.

To express "sealed, but may read the op's operands," hipsr defines a standalone trait, IsolatedFromAboveButAllowOperands. MLIR's built-in passes don't know about it — they act only on the standard sealing marker, IsIsolatedFromAbove. So a standard pass sees no boundary on a hipsr region and treats it as an ordinary, non-isolated region it can freely reach into and rewrite.

The result is a mismatch: the passes move values across a boundary they can't see, so the IR they produce breaks the region's sealing rule — it is illegal. This issue shows three concrete cases, each reproducible on the merged hipsr.cast op, and asks how the pipeline plans to stay correct when standard MLIR passes run.

The failure is a property of the IsolatedFromAboveButAllowOperands trait, not of hipsr.cast specifically: every Hipsr_DpsOp carries that trait, and the same -cse / -canonicalize errors reproduce on hipsr.pool_domain (the other merged op that carries it). The examples below use the merged hipsr.cast + hipsr.placeholder ops (both on main) only because their shape-region examples are the clearest.

(Shape population — populateShapeRegion — emits shape.shape_of + shape.get_extent; the shape dialect isn't registered in hip-mlir-opt, so the examples hand-write the equivalent with tensor.dim — same meaning: read an input's size, same trait, same downstream behavior.)

Environment

• hip-mlir-opt built from current main (which has the merged hipsr.cast, hipsr.placeholder, and hipsr.pool_domain ops, plus the empty-region allowance on Hipsr_DpsOp).

---

Case 1 — -cse pulls an outside value into the region, producing illegal IR

CSE (common-subexpression elimination) removes duplicate computations. Here the function computes a size (tensor.dim %input) once outside the op and once inside the op's region. Because CSE doesn't see the region as sealed, it merges the two — rewriting the inside copy to point at the outside one. Now the region reads a value from outside itself, which the sealing rule forbids.

func.func @cast_cse_capture(%input: tensor<?x8xf32>) -> (index, tensor<?x8xf16>) {
  %c0 = arith.constant 0 : index
  %p  = tensor.dim %input, %c0 : tensor<?x8xf32>        // computed once, outside the op
  %init = hipsr.placeholder : tensor<?x8xf16>
  %0 = hipsr.cast ins(%input : tensor<?x8xf32>) outs(%init : tensor<?x8xf16>)
                  -> tensor<?x8xf16> shape_region {
    %c0i = arith.constant 0 : index
    %d0  = tensor.dim %input, %c0i : tensor<?x8xf32>    // the same size, computed inside
    %c8  = arith.constant 8 : index
    hipsr.shape_yield (%d0, %c8) : [f16]
  }
  return %p, %0 : index, tensor<?x8xf16>
}

$ hip-mlir-opt -cse cast_cse_capture.mlir
error: 'hipsr.shape_yield' op using value defined outside the region
    hipsr.shape_yield (%d0, %c8) : [f16]
    ^
note: may only use values defined in its regions or the op's operands

This isn't a contrived input. Everything above is legal IR, and from CSE's point of view the rewrite — merging a duplicate into the value that dominates it — is a routine, semantics-preserving step. But the result violates the region's sealing invariant, so the IR is now illegal; the verifier only reports what the pass already produced. Any graph in which a dim computed inside a region also appears at a point that dominates the op (main-graph shape arithmetic, or a dim an earlier pass reified out) hits this under a plain -cse.

---

Case 2 — -canonicalize moves a folded constant out of the region, producing illegal IR

When canonicalization folds a constant (here 2 * 4 = 8), it places the new constant at the top of the nearest sealed region so it can be shared. It checks for the standard sealing marker, doesn't find it on the hipsr op, and so keeps going up and puts the constant in the enclosing function — outside the region. The region's shape_yield then reads a value that now lives outside it.

func.func @cast_fold_hoist(%input: tensor<8xf32>) -> tensor<8xf16> {
  %init = hipsr.placeholder : tensor<8xf16>
  %0 = hipsr.cast ins(%input : tensor<8xf32>) outs(%init : tensor<8xf16>)
                  -> tensor<8xf16> shape_region {
    %c2 = arith.constant 2 : index
    %c4 = arith.constant 4 : index
    %f  = arith.muli %c2, %c4 : index   // folds to 8
    hipsr.shape_yield (%f) : [f16]
  }
  return %0 : tensor<8xf16>
}

$ hip-mlir-opt -canonicalize cast_fold_hoist.mlir
error: 'hipsr.shape_yield' op using value defined outside the region
    hipsr.shape_yield (%f) : [f16]
    ^
note: may only use values defined in its regions or the op's operands

Even the most basic canonicalization — folding a constant — breaks the sealing rule.

---

Case 3 — loop-invariant code motion can't reach into the region (no error, just a missed optimization)

LICM moves a whole operation out of a loop when its result doesn't change between iterations. It does not look inside an operation's region to move a piece of it out. Here the cast as a whole does change each iteration (its running buffer is loop-carried), but the size it computes inside its region (tensor.dim %input) does not. LICM can't move the whole op, and won't reach into the region to lift just the size — so that size is recomputed every iteration. (Threading the loop-carried buffer through the init is an illustrative way to make the op loop-variant while its region's shape stays invariant; the general point is that whenever those two hold, the region's shape math can't be lifted out.)

func.func @cast_licm(%input: tensor<?x8xf32>, %seed: tensor<?x8xf16>,
                     %lb: index, %ub: index, %step: index) -> tensor<?x8xf16> {
  %r = scf.for %i = %lb to %ub step %step iter_args(%acc = %seed) -> (tensor<?x8xf16>) {
    %0 = hipsr.cast ins(%input : tensor<?x8xf32>) outs(%acc : tensor<?x8xf16>)
                    -> tensor<?x8xf16> shape_region {
      %c0 = arith.constant 0 : index
      %d0 = tensor.dim %input, %c0 : tensor<?x8xf32>   // same every iteration, but sealed in
      %c8 = arith.constant 8 : index
      hipsr.shape_yield (%d0, %c8) : [f16]
    }
    scf.yield %0 : tensor<?x8xf16>
  }
  return %r : tensor<?x8xf16>
}

After -loop-invariant-code-motion, tensor.dim %input stays inside the region, inside the loop. Written the ordinary way (the size as a plain value in the loop body, not sealed in a region), the same pass lifts it above the loop:

// ordinary version, after -loop-invariant-code-motion: the size is lifted above the loop
%dim = tensor.dim %input, %c0 : tensor<?x8xf32>
%0 = scf.for ... { ... }

LICM is also an enabling pass: a hoisted invariant becomes visible to CSE, constant folding, and further hoisting. So sealing the size in the region doesn't just cost one recompute per iteration — it forecloses the follow-on optimizations the hoist would have unlocked.

What this issue is not saying

• It is not saying the pipeline errors out end-to-end today — that depends on whether any standard MLIR pass runs while a region is still populated.
• Case 3 is recovered if the region is reified into ordinary values before those passes run — but that reification is the ordinary value-based approach.

Open questions

1. A standard MLIR pass produces the illegal IR before the verifier can reject it, so staying correct means ensuring no such pass ever runs while a region is populated. Is there a way to guarantee that — enforceable for all future passes, front-ends, and outside contributors — and if so, how?
2. If the shape is instead reified into ordinary value-graph SSA before those passes run (which is what makes it safe), what does the region add over keeping the shape as ordinary SSA from the start?

[fhanuman]

Notify -- @amd-bqiao, @edelaye, @wcy123, @amd-mingw, @sapareek, @qianglin-amd

[amd-mingw]

Hi @fhanuman @edelaye,
Thank you for your feedback, this is truly positive and valuable input. There are mainly two issues in this feedback: one is a specific problem, and the other is a methodological point.

1. The specific problem is the custom trait IsolatedFromAboveButAllowOperands, which can be solved directly by using the standard trait IsolatedFromAbove, refactor(hipsr): isolate shape regions and clean up ShapeRegionInterface #519 is working on it.
2. The methodological point requires that the input and output MLIR throughout the entire pipeline must be standard and legal MLIR, it's a good point.

About the specific problem(custom trait IsolatedFromAboveButAllowOperands) and #519 , detail analysis and solution is:
"IsolatedFromAbove only prevents direct SSA capture from containing scopes. Values flow through block arguments."

From MLIR Language Reference:

By default, operations inside a region can reference values defined outside of the region whenever it would have been legal for operands of the enclosing operation to reference those values, but this can be restricted using traits, such as OpTrait::IsolatedFromAbove.

How Standard Dialects Handle This:

• func.func - IsolatedFromAbove, function arguments are block arguments
• scf.for - IsolatedFromAbove, loop IV and iter_args are block arguments
• omp.target - IsolatedFromAbove, captured values passed through host_eval operands → block arguments
• linalg.generic - IsolatedFromAbove, input/output tensor elements passed as block arguments

All follow the pattern: IsolatedFromAbove + block arguments for external values.

Proposed Solution High-Level Design(#519 )

• Use standard IsolatedFromAbove trait (not custom variant)
• Pass operand values through block arguments following MLIR convention
• Provide named accessor helpers for readability
• Remove IsolatedFromAboveButAllowOperands trait entirely

cc @wcy123

[fhanuman]

Hi @amd-mingw, thanks for the clarification. @amd-bqiao, @edelaye, and I met to discuss the proposed change. It resolves the illegal-IR issue, but while the shape remains inside the isolated region, stock CSE, canonicalization, and LICM cannot share or hoist its computations across the boundary. I have captured our findings and questions in the issue below—please take a look when you get a chance:

#525

One clarification on the dialect examples in your list: func.func and omp.target are IsolatedFromAbove, but scf.for and linalg.generic are not. Their bodies may capture values directly from the enclosing scope; only the IV/iter_args or payload element values are block arguments. Annotating those two as IsolatedFromAbove would make currently legal IR that captures outer values illegal.

IsolatedFromAbove is used on independently processed units such as func.func, omp.target, and gpu.module. It forbids implicit SSA captures and is required for operations that anchor nested pass pipelines; omp.target additionally implements the OpenMP dialect's outlining interface (OutlineableOpenMPOpInterface). This differs from local computation regions such as scf.for and linalg.generic. I am not saying it cannot be used for region-based shape inference—only that doing so brings the cross-boundary optimization cost described in the issue.