[RFC] HIP EP result shapes, dynamic pool packing, and host visibility

[fhanuman]

Summary

The HIP EP lowers ONNX graphs through MLIR to GPU code. It creates dynamic tensor destinations before bufferization and packs temporary buffers into session-owned GPU pools.

The shipping path has three separable problems:

1. ONNX→HIP conversion and reifyResultShapes have independently computed the same result shapes and have disagreed for valid broadcast, Gemm, and MatMul inputs.
2. For runtime-sized allocations, PoolAllocs may reserve separate slabs for buffers whose lifetimes do not overlap because their size expressions are different SSA values.
3. Some output sizes and control values come from tensors that may not yet be available to host code. Bare host extracts are unsafe, while encoding synchronized backend reads directly in normalized IR selects the visibility mechanism too early.

Result-shape agreement and safe host visibility are correctness work. Lifetime packing is a separate optimization. Their generated SSA can interact, but an improvement in one area does not imply fewer pools or lower memory in another.

Each fix stays within the existing hip.* path:

• Fix 1 — shared result-shape implementation: use one shape helper for destination construction and reification.
• Fix 2 — lifetime-only dynamic packing: change dynamic allocation grouping and slab sizing inside PoolAllocs.
• Fix 3 — scalar extraction and host visibility: propose separating which tensor value is needed from how that value becomes available to host code. ONNX→HIP conversion expresses scalar extraction; the exact backend lowering remains TBD.

No new dialect is proposed.

Review requested

Three shipping-path contracts require review:

1. Shared result shapes. Converter destination construction and reifyResultShapes call one OpFoldResult shape implementation. Equivalent dynamic-extent SSA is reused or canonicalized before buffer packing.
2. Lifetime-only packing. Runtime-sized allocations may share an aligned pool offset only when their lifetimes are disjoint. Requested capacity and retained high-water are measured against unchanged packing.
3. Scalar extraction and host visibility. Decide whether normalized ONNX→HIP IR should describe only the tensor element being read, leaving backend lowering responsible for making that value safely available to host shape or control logic.

Scope and non-goals

This RFC changes the shipping hip.* path. It does not:

• introduce a new shape IR or replace InferTypeOpInterface, ReifyRankedShapedTypeOpInterface, or DestinationStyleOpInterface;
• add a general scheduling or synchronization abstraction;
• change graph output allocation, inference_compute, or control-flow ABIs.

Background: what exists today

Shipping pipeline

ONNX
→ cleanup and control-flow outlining
→ ONNX-to-HIP conversion and shape refinement
→ One-Shot Bufferize and output-allocation rewriting
→ memref cleanup, host-scalar materialization, and dimension resolution
→ PoolAllocs
→ allocation and external-constant lowering
→ HIP-to-LLVM lowering and operation-state setup
→ generate-interface

lib/Dialect/Transforms/Pipelines.cpp is the source of truth for exact pass order. This RFC names individual passes only where their placement affects a contract.

Ownership-based deallocation is absent from the shipping pipeline. Every temporary allocation must enter a session-owned pool; graph outputs become hip.alloc_output.

hip-use-output-allocator runs before pooling. hip.alloc_output is distinct from memref.alloc, and PoolAllocs rewrites only memref.alloc. The generated interface unconditionally uses the two-argument inference_compute(state, inputs) ABI.

The path uses standard upstream MLIR concepts:

• DestinationStyleOpInterface for tensor destinations;
• InferTypeOpInterface for result types tied to destinations;
• ReifyRankedShapedTypeOpInterface for static attributes and dynamic extent SSA;
• ordinary tensor.dim, memref.dim, arith, and control-flow SSA;
• explicit side-effecting HIP operations for GPU→host visibility;
• post-bufferization liveness and arena-style allocation packing.

Existing result-shape flow

MLIR represents each reified dimension as an OpFoldResult: either a static integer attribute or a runtime SSA value.

ONNX→HIP conversion must create a destination before constructing a HIP destination-passing-style (DPS) op:

%destination_extent = tensor.dim %input, %c0
%init = tensor.empty(%destination_extent) : tensor<?xf32>
%result = hip.op ... outs(%init : tensor<?xf32>) -> tensor<?xf32>
%used_dim = tensor.dim %result, %c0 : tensor<?xf32>

ReifyRankedShapedTypeOpInterface is then queried out-of-band; it is not a second operation in the printed IR:

reifyResultShapes(%result) -> [%reified_extent]

Shape refinement and tensor.dim folding use %reified_extent. The contract is:

%destination_extent and %reified_extent implement the same shape function

If the two expressions implement different shape rules, the destination and the shape observed by consumers disagree. Problem 1 — one result shape is computed twice gives a concrete Max example.

Existing PoolAllocs flow

After bufferization, PoolAllocs:

1. computes allocation liveness intervals;
2. partitions allocations into dominance-feasible domains;
3. groups allocations and assigns aligned offsets within each domain;
4. replaces allocations with views into grow-on-demand runtime pools.

The shipping pass requires a single-block function. Each lifetime runs from the memref.alloc to the latest transitive use through forward aliases; explicit memref.dealloc operations are removed after pooling.

A simplified single-domain transformation is:

// Before PoolAllocs.
%a = memref.alloc(%m) : memref<?xf16>
%b = memref.alloc(%n) : memref<?xf16>

// After PoolAllocs. f16 footprints are aligned to 256 bytes.
%a_bytes = arith.muli %m, %c2 : index
%a_ceil = arith.addi %a_bytes, %c255 : index
%a_chunks = arith.divui %a_ceil, %c256 : index
%a_aligned = arith.muli %a_chunks, %c256 : index

%b_bytes = arith.muli %n, %c2 : index
%b_ceil = arith.addi %b_bytes, %c255 : index
%b_chunks = arith.divui %b_ceil, %c256 : index
%b_aligned = arith.muli %b_chunks, %c256 : index

%pool_size = arith.addi %a_aligned, %b_aligned : index
%pool = hip.get_pool(%ctx, %pool_size)
    {domain_id = 0 : i64} : memref<?xi8>
%a_view = memref.view %pool[%c0][%m]
    : memref<?xi8> to memref<?xf16>
%b_view = memref.view %pool[%a_aligned][%n]
    : memref<?xi8> to memref<?xf16>

The exact offsets depend on liveness and grouping. Problem 2 — dynamic buffers are grouped before lifetime is considered shows the case where lifetime-disjoint dynamic buffers should reuse one offset instead of using the summed layout above.

A dominance domain is a set of allocations whose size expressions can be computed at one legal insertion point. A synchronized readback in the middle of a function can therefore force later allocations into a separate domain.

Dynamic allocation sizes are SSA expressions. Existing grouping can treat different expressions as separate units even when the corresponding allocations are lifetime-disjoint.

Current grouping does not allow different dynamic-size SSA groups to share one slab.

Existing runtime-value readback

Metadata such as tensor.dim %input is host-available descriptor information. A scalar value written by a GPU kernel is different: SSA dominance alone does not make device bytes host-visible.

The HIP dialect provides:

%v = hip.readback_scalar(%ctx, %device_scalar : tensor<i64>) -> i64
%n = hip.readback_dim(%ctx, %device_count : tensor<i32>) -> index

readback_dim is the specialized i32-count-to-index form. readback_scalar preserves the source element type.

A descriptor extent is already available to host shape arithmetic:

%n = tensor.dim %input, %c0 : tensor<?xf32>
%init = tensor.empty(%n) : tensor<?xf32>

hip.nonzero is different because the number of non-zero elements depends on tensor payload values. It returns a capacity-sized indices tensor and a scalar containing the valid count:

%dim0 = tensor.dim %input, %c0 : tensor<?x?xi64>
%dim1 = tensor.dim %input, %c1 : tensor<?x?xi64>
%capacity = arith.muli %dim0, %dim1 : index
%indices_init = tensor.empty(%capacity) : tensor<2x?xi64>
%count_init = tensor.empty() : tensor<i32>
%nz:2 = hip.nonzero(%ctx)
    ins(%input : tensor<?x?xi64>)
    outs(%indices_init, %count_init : tensor<2x?xi64>, tensor<i32>)
    {input_data_type = 4 : i64}
    : tensor<2x?xi64>, tensor<i32>
%n = hip.readback_dim(%ctx, %nz#1 : tensor<i32>) -> index
%valid = tensor.extract_slice %nz#0[0, 0] [2, %n] [1, 1]
    : tensor<2x?xi64> to tensor<2x?xi64>

Here %n is the runtime count used to trim the capacity-sized indices tensor to its valid shape. Problem 3 asks whether normalized conversion should describe that scalar access without selecting the backend visibility mechanism.

Problem 1: one result shape is computed twice

Before the proposed fix, shipping ONNX→HIP conversion computes dynamic extents to build a DPS destination. Later, reifyResultShapes computes the extents again for shape refinement and tensor.dim folding. These two paths use different code.

The required invariant is simple:

Destination construction and result-shape reification must implement the same ONNX shape function.

Minimal mismatch

Consider a fully dynamic Max:

func.func @max_dynamic(
    %lhs: tensor<?x?xf32>,
    %rhs: tensor<?x?xf32>) -> tensor<?x?xf32> {
  %result = "onnx.Max"(%lhs, %rhs)
      : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
  return %result : tensor<?x?xf32>
}

For one valid runtime input:

lhs shape = [4, 1]
rhs shape = [4, 8]
result    = [4, 8]

The old converter chose dimensions from one operand:

%lhs_m = tensor.dim %lhs, %c0
%lhs_n = tensor.dim %lhs, %c1
%init = tensor.empty(%lhs_m, %lhs_n) : tensor<?x?xf32>
%result = hip.max ... outs(%init : tensor<?x?xf32>)

At runtime, this destination is [4, 1].

The old reifier used broadcast semantics:

%lhs_m = tensor.dim %lhs, %c0
%rhs_m = tensor.dim %rhs, %c0
%m_lhs_is_one = arith.cmpi eq, %lhs_m, %c1 : index
%m = arith.select %m_lhs_is_one, %rhs_m, %lhs_m : index

%lhs_n = tensor.dim %lhs, %c1
%rhs_n = tensor.dim %rhs, %c1
%n_lhs_is_one = arith.cmpi eq, %lhs_n, %c1 : index
%n = arith.select %n_lhs_is_one, %rhs_n, %lhs_n : index
// reifyResultShapes returns [%m, %n].

For the same input, reification returns [4, 8]. One HIP operation therefore has a [4, 1] destination but a [4, 8] reified result shape. The destination is wrong before bufferization or PoolAllocs runs.

Here converter destination construction is wrong; reifyResultShapes yields the ONNX broadcast shape. Gemm and MatMul have analogous converter/reifier dimension-source mismatches covered by Fix 1 — shared result-shape implementation.

For ONNX-compatible extents, the shared broadcast rule must handle (1, N), (N, 1), (N, N), and zero extents. It cannot use integer maximum because ONNX broadcasting of 0 and 1 produces 0.

Required fix

Both paths call the same helper:

%lhs_m = tensor.dim %lhs, %c0
%rhs_m = tensor.dim %rhs, %c0
%m_lhs_is_one = arith.cmpi eq, %lhs_m, %c1 : index
%m = arith.select %m_lhs_is_one, %rhs_m, %lhs_m : index

%lhs_n = tensor.dim %lhs, %c1
%rhs_n = tensor.dim %rhs, %c1
%n_lhs_is_one = arith.cmpi eq, %lhs_n, %c1 : index
%n = arith.select %n_lhs_is_one, %rhs_n, %lhs_n : index

%init = tensor.empty(%m, %n) : tensor<?x?xf32>
%result = hip.max ... outs(%init : tensor<?x?xf32>)

reifyResultShapes calls the same helper and returns semantically equivalent [%m, %n] dimensions. Fix 1 — shared result-shape implementation also defines the structural requirement for equivalent extent SSA materialized by separate calls.

Problem 2: dynamic buffers are grouped before lifetime is considered

The required invariant is:

Allocations may share a pool offset only when their lifetimes are disjoint; different dynamic-size SSA values alone must not prevent that sharing.

This issue is specific to dynamic shapes. Static allocations have constant byte footprints and do not require runtime-size SSA grouping.

Consider two dynamic allocations in one PoolAllocs domain:

func.func @lifetime_disjoint(
    %ctx: !hip.context, %m: index, %n: index,
    %xa: memref<?xf16>, %ya: memref<?xf16>,
    %xb: memref<?xf16>, %yb: memref<?xf16>) {
  %a = memref.alloc(%m) : memref<?xf16>
  hip.mul(%ctx)
      ins(%xa, %ya : memref<?xf16>, memref<?xf16>)
      outs(%a : memref<?xf16>)
  // %a's last use is the hip.mul above.

  %b = memref.alloc(%n) : memref<?xf16>
  hip.mul(%ctx)
      ins(%xb, %yb : memref<?xf16>, memref<?xf16>)
      outs(%b : memref<?xf16>)
  return
}

Their lifetimes do not overlap, but their size SSA values differ. Grouping by the size expression before exploiting lifetime can reserve separate storage:

%a_bytes_aligned = ...
%b_bytes_aligned = ...
%reserved = arith.addi %a_bytes_aligned, %b_bytes_aligned : index
%pool = hip.get_pool(%ctx, %reserved) {domain_id = 0 : i64}
%a = memref.view %pool[%c0][%m] : memref<?xi8> to memref<?xf16>
%b = memref.view %pool[%a_bytes_aligned][%n]
    : memref<?xi8> to memref<?xf16>

Because the lifetimes are disjoint, both views can start at offset zero and the required peak is the aligned runtime maximum:

%required = arith.maxui %a_bytes_aligned, %b_bytes_aligned : index
%pool = hip.get_pool(%ctx, %required) {domain_id = 0 : i64}
%a = memref.view %pool[%c0][%m] : memref<?xi8> to memref<?xf16>
%b = memref.view %pool[%c0][%n] : memref<?xi8> to memref<?xf16>

Alignment is part of each footprint before the maximum is computed. The shipping pipeline does not insert per-allocation deallocations; PoolAllocs derives interval ends from the last transitive use.

Problem 3: scalar extraction and host visibility are mixed

Assume %shape contains the runtime size of an output tensor. Host code needs that size before it can create the output:

%length = ...
%extent = arith.index_castui %length : i64 to index
%output = tensor.empty(%extent) : tensor<?xf32>

The current IR tells ONNX→HIP conversion how to obtain %length:

%length = hip.readback_scalar(
    %ctx, %shape : tensor<i64>) -> i64
%extent = arith.index_castui %length : i64 to index
%output = tensor.empty(%extent) : tensor<?xf32>

hip.readback_scalar does more than identify the value. It selects a backend-specific way to make that value available to host code.

Normalized IR only needs to describe which value is required:

%length = tensor.extract %shape[] : tensor<i64>
%extent = arith.index_castui %length : i64 to index
%output = tensor.empty(%extent) : tensor<?xf32>

Both examples have the same model meaning: read the output length and use it to create %output.

The difference is who decides how %length becomes available:

ONNX→HIP conversion → identifies the scalar value
backend lowering     → loads it directly, synchronizes first, or reports an error

This keeps backend-specific visibility details out of normalized IR.

Proposed fixes

Fix 1: shared result-shape implementation

Shared contract

HipShapeUtils owns shared broadcast, Gemm, and MatMul helpers. They return OpFoldResult dimensions so static extents remain attributes and dynamic extents remain SSA values. Converter destination construction applies the rank and static constraints of the imported ONNX result type; reification calls the same helper instead of recomputing the semantics.

The converter path is:

FailureOr<SmallVector<OpFoldResult>> dims =
    reifyBroadcastResultShape(builder, loc, operands, emitError);
FailureOr<Value> init =
    createEmptyTensorFromReifiedShape(builder, loc, resultType, *dims);

The op's ReifyRankedShapedTypeOpInterface implementation calls the same shape helper. Static attributes and dynamic SSA values therefore follow one ONNX shape function.

Together, these shared helpers cover right-aligned broadcast, variadic Max/Min, transpose-aware Gemm, and batched MatMul, including rank-zero results. Equivalent dynamic extents are reused or canonicalized before bufferization to avoid duplicate allocation-size SSA.

A draft implementation is available in PR #581.

Fix 2: lifetime-only dynamic packing

Within each dominance domain:

1. begin with each runtime-sized allocation in its own candidate group;
2. merge groups only when every member lifetime is disjoint;
3. place merged members at one aligned offset and size the shared slab with the runtime maximum of their aligned footprints.

The merged group's liveness is the union of its members. Any overlapping member pair prevents the merge. Size/max expressions are emitted at the domain anchor after their operand definitions and before every replaced allocation.

A draft implementation is available in PR #570.

Fix 3: proposed separation of scalar extraction and host visibility

ONNX→HIP conversion uses ordinary scalar SSA:

• compile-time values may be folded;
• descriptor extents use tensor.dim or memref.dim;
• tensor payload values use tensor.extract.

Before host code uses a payload scalar, backend-specific lowering must either prove the value is already available, make it available safely, or report an error.

The exact lowering design is TBD. A draft implementation is needed to determine:

• where the visibility decision should run;
• how the compiler tracks whether a value is ready for host use;
• how synchronized loads are kept out of the ordinary load-lowering path;
• how copy or synchronization failures propagate to dependent operations.

Draft implementations

• Fix 1: PR #581
• Fix 2: PR #570
• Fix 3: implementation TBD

References

• mlir/include/mlir/Interfaces/InferTypeOpInterface.h
• mlir/include/mlir/Interfaces/SideEffectInterfaces.td
• include/hip/Dialect/IR/HipOps.td
• include/hip/Dialect/IR/HipBufferize.h
• lib/Dialect/Transforms/PoolAllocs.cpp
• docs/design/hip-shape-inference.md
• docs/design/pool-allocs-memory-planning.md
• PR #581
• PR #570