Skip to content

fix(hip): share Gather and tensor-shape contracts #644

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
fhanuman wants to merge 1 commit into from
Draft

fix(hip): share Gather and tensor-shape contracts #644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions docs/design/compiler-runtime-contract.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,16 @@ batch/stride assumptions plus a true axis-validity bit. HIP artifacts built
against an earlier checked MatMul signature and artifacts built against the old
one-K Gemm signature must be recompiled; existing Hipsr MatMul artifacts remain
valid.

GatherND has another generated-code-only contract. `wrap_gather_nd` receives
the data, indices, and output pointers plus their host-side i64 shape arrays,
ranks, `batch_dims`, and the data element type. It carries no indices element
width. Both the wrapper and custom kernel interpret `indices` as an
`int64_t *`, so ONNX conversion, HIP verification and reification, and
HIP-to-LLVM lowering all require i64 indices before emitting destination-shape
IR or a runtime call. Supporting i32 indices would require an explicit runtime
ABI extension rather than reusing this call.

---

## Consumers
Expand Down
49 changes: 45 additions & 4 deletions docs/design/hip-shape-inference.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,11 +87,13 @@ Choose the smallest mechanism that matches the operation's semantics:
| Shape contract | Mechanism |
|---|---|
| Result shape equals DPS init shape, including most multi-result DPS ops | Shared `HipDpsOpInterface` default |
| Result shape equals a named input | `reifyElementwiseSameShape` or a small dedicated thunk |
| Result shape equals a named input | `Hip_DpsOp_SameShape`, `reifyElementwiseSameShape`, and `verifySameShapeDpsOp` |
| NumPy-style broadcast | `Hip_DpsOp_Broadcast`, `reifyBroadcastResultShape`, and the shared converter bridge |
| Reduction with constant axes/keepdims | `Hip_DpsOp_Reduction` and reduction helpers |
| Permutation | `reifyTransposeByPerm` |
| Gather/GatherND/GatherElements | Gather-specific helpers or thunks |
| Permutation | Shared pure `inferTransposeShape` plus `reifyTransposeByPerm` and verifier |
| Gather/GatherND | Shared pure Gather/GatherND helpers plus reify/verifier thunks; dynamic GatherND tuple width falls back to outs |
| GatherElements | `Hip_DpsOp_SameShape` with `indices` as its named source |
| GatherBlockQuantized | Shared logical-dequantized Gather rule; packed int4 expands the surviving quantize axis |
| OneHot, Compress, TopK | Dedicated reification thunks |
| Pad | Exact affine input-dimension arithmetic for constant/stamped pads; runtime pads use converter readback and reify from outs |
| Slice | Shared static/SSA normalization; runtime i32/i64 controls use one grouped readback and exact extents |
Expand All @@ -104,7 +106,7 @@ Choose the smallest mechanism that matches the operation's semantics:
| Rank-4 NCHW ConvTranspose | Shared ONNX formula used by converter, reification, and verifier |
| GlobalPool | N/C from input and unit spatial extents, shared by converter, reification, and verifier |
| CausalConvWithState | Runtime-supported 1D output/state formulas from input and depthwise kernel |
| Resize | DPS-init shape, with semantic validity handled by conversion |
| Resize | N/C from input plus static spatial extents from the imported output template |
| Runtime-dependent count, such as NonZero or Compress | DPS-init shape; unresolved dimensions remain dynamic |

Shared declarations live in `HipShapeUtils.h`; common implementation lives in
Expand Down Expand Up @@ -220,6 +222,29 @@ produces 0, not 1. Variadic Max/Min share one
`lowerVariadicBroadcastChain` helper that derives every pairwise intermediate
type from this shared broadcast shape.

GatherND accepts only i64 indices at the HIP boundary. Its runtime call carries
no index-width argument, and the custom kernel reads the indices pointer as
`int64_t *`; conversion therefore rejects i32 before emitting shape SSA or a
destination, while verification, reification, and lowering enforce the same
contract defensively. A static trailing tuple width uses the shared exact shape
rule. An already-formed HIP op whose i64 tuple width is dynamic remains legal:
its reifier cannot know the result rank from the operands and falls back to the
DPS destination shape.

GatherBlockQuantized applies Gather to the logical dequantized data shape, not
the physical byte-storage shape. For the HIP op's byte-packed 4-bit storage,
`logical_data[quantize_axis] = 2 * data[quantize_axis]`; all other extents are
unchanged. The result is
`logical_data[:gather_axis] ++ indices.shape ++
logical_data[gather_axis+1:]`. If Gather removes the quantize axis itself, no
result extent is doubled. The shared helper also validates the statically known
block grid (`scales[quantize_axis] =
ceil(logical_data[quantize_axis] / block_size)`), non-quantized extents,
optional zero-point packing, axes, ranks, and runtime-supported attributes.
Converter destination construction, op reification, and static verification all
call this rule. The runtime receives the physical data descriptor and expands
its quantize-axis extent before launching the logical-element kernel.

Forward Conv and Pool share one validated spatial-window primitive:
`floor((input + pads - effectiveKernel) / stride) + 1`; Pool selects signed
ceil division when `ceil_mode = 1`. After that ceil calculation, a positive
Expand Down Expand Up @@ -485,6 +510,15 @@ be refined by a proven static extent. An imported static dimension requires an
equal static inferred extent; unknown inference never inherits an importer
constraint. A failed pattern therefore leaves the IR unchanged.

Resize is semantic rather than payload-dependent at the HIP level. ONNX
`sizes`/`scales` have already been resolved by the importer and are not operands
of `hip.resize`; the runtime recovers each scale from the input/output
descriptors. The shared rule therefore takes N/C from the input and the spatial
extents from the imported output type, requiring every spatial extent to be
static. A dynamic input N/C extent requires a dynamic output template extent;
a static input extent may refine a dynamic template. This deliberately does not
claim dynamic spatial support without carrying the payload.

## Pre-conversion loop-body rank inference

`--hip-infer-loop-body-shapes` is a narrow pre-conversion backstop. It runs after loop outlining and before ONNX-to-HIP conversion to establish rank for unranked loop-carried values that would otherwise block conversion; it is not the general HIP shape-inference mechanism.
Expand Down Expand Up @@ -524,6 +558,13 @@ Primary regression coverage:
| `test/lit/Dialect/hip-broadcast-shape-verifier.mlir` | Generated NumPy broadcast shape verification |
| `test/lit/Dialect/hip-reduction-shape-verifier.mlir` | Generated reduction shape and axes verification |
| `test/lit/Conversion/onnx-to-hip/test_reshape_shape_provenance.mlir` | Proven host Reshape shapes, dataflow joins/shared producers, unknown-payload fallback, and `-1` handling |
| `test/lit/Conversion/onnx-to-hip/test_gather_nd_invalid.mlir` | Mutation-free rejection of i32 GatherND indices |
| `test/lit/Dialect/hip-gather-nd-reify-shapes.mlir` | Static tuple-width reification and dynamic tuple-width destination fallback |
| `test/lit/Dialect/hip-gather-nd-shape-verifier.mlir` | GatherND i64 ABI and static/dynamic tuple-width verification |
| `test/lit/Dialect/hip-gather-block-quantized-reify-shapes.mlir` | Packed-int4 logical Gather shape reification |
| `test/lit/Dialect/hip-gather-block-quantized-shape-verifier.mlir` | GatherBlockQuantized logical shape and runtime-contract validation |
| `test/lit/Dialect/hip-resize-reify-shapes.mlir` | Dynamic N/C and static spatial Resize reification |
| `test/lit/Dialect/hip-resize-shape-verifier.mlir` | Resize semantic destination validation |
| `test/lit/Dialect/hip-matmul-reify-shapes.mlir` | Per-op reification through `--resolve-shaped-type-result-dims` |
| `test/lit/Dialect/hip-matmul-shape-verifier.mlir` | Static MatMul shape validation |
| `test/lit/Conversion/onnx-to-hip/test_reduce_sum.mlir` | Reduction destinations, including the non-positional `keepdims = 0` dimension mapping |
Expand Down
2 changes: 1 addition & 1 deletion docs/supported-operations.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,7 @@ The conversion registrations in `lib/Conversion/OnnxToHip/OnnxToHip.cpp` and the
| MaxPool | Custom HIP kernel |
| AveragePool | Custom HIP kernel |
| LpPool | Custom HIP kernel |
| Resize | Custom HIP kernel |
| Resize | Trailing 1D/2D/3D spatial resize with static output spatial extents; dynamic N/C pass through from input |
| GlobalAveragePool | Custom HIP kernel |
| GlobalMaxPool | Custom HIP kernel |
| GlobalLpPool | Custom HIP kernel |
Expand Down
97 changes: 58 additions & 39 deletions include/hip/Dialect/IR/HipOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1203,7 +1203,9 @@ def Hip_QMoEOp : Hip_DpsOp<"qmoe",
}];
}

def Hip_GatherBlockQuantizedOp : Hip_DpsOp<"gather_block_quantized"> {
def Hip_GatherBlockQuantizedOp :
Hip_DpsOp_Semantic<"gather_block_quantized", /*traits=*/[],
/*outsAccessor=*/"Output", /*autoReify=*/0> {
let summary = "Gather + block-wise dequantize "
"(com.microsoft GatherBlockQuantized)";
let description = [{
Expand All @@ -1212,17 +1214,22 @@ def Hip_GatherBlockQuantizedOp : Hip_DpsOp<"gather_block_quantized"> {
integer `indices`, and dequantize the gathered rows using per-block
`scales` and optional `zero_points`.

`data` is a constant tensor of rank `r >= 1` quantized along
`data` is a constant tensor of rank `r > 1` quantized along
`quantize_axis` with the given `block_size`. For 4-bit, two values are
packed per byte (low nibble first, high nibble second); for 8-bit one
value per byte. `scales` (and `zero_points`, when present) share the
layout of `data` except along `quantize_axis`, where their extent is
`data.shape[quantize_axis] / block_size`. When `zero_points` is absent,
the default is `0` for int4/uint4 and `2^(bits-1)` for uint8.

Output rank is `q + (r - 1)` where `q = rank(indices)`. Element type of
`output` matches `scales` (T2: f32 / f16 / bf16). For uint8 `data`,
`gather_axis` must be 0.
`ceil(logical_data.shape[quantize_axis] / block_size)`. When
`zero_points` is absent, the default is `0` for signed storage and
`2^(bits-1)` for unsigned storage.

The result is Gather over the logical dequantized data shape:
`logical_data[:gather_axis] ++ indices.shape ++
logical_data[gather_axis+1:]`. For byte-packed 4-bit storage,
`logical_data[quantize_axis] = data.shape[quantize_axis] * 2`; no extent is
doubled when Gather removes that axis. Output rank is `q + (r - 1)` where
`q = rank(indices)`. Element type of `output` matches `scales`
(T2: f32 / f16 / bf16). For uint8 `data`, `gather_axis` must be 0.

Optional unit attribute `unsigned_quant_storage`: set when
convert-onnx-to-hip identifies UINT4/UINT8
Expand All @@ -1243,7 +1250,7 @@ def Hip_GatherBlockQuantizedOp : Hip_DpsOp<"gather_block_quantized"> {
memref<8xi64, 1>,
memref<2048x12xf16, 1>)
zero_points(%zp : memref<2048x12xui8, 1>)
outs(%out : memref<8x96xf16, 1>)
outs(%out : memref<8x192xf16, 1>)
{bits = 4, block_size = 16, gather_axis = 0, quantize_axis = 1}
```
}];
Expand All @@ -1269,6 +1276,7 @@ def Hip_GatherBlockQuantizedOp : Hip_DpsOp<"gather_block_quantized"> {
`outs` `(` $output `:` type($output) `)`
attr-dict (`:` type($result_tensors)^)?
}];
let hasVerifier = 1;
}

// ===== MIOpen ops ============================================================
Expand Down Expand Up @@ -1791,11 +1799,12 @@ def Hip_MiopenSoftmaxOp : Hip_DpsOp<"miopen.softmax",

// ===== Custom HIP kernel ops (no library equivalent) =========================

def Hip_TransposeOp : Hip_DpsOp<"transpose", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
def Hip_TransposeOp :
Hip_DpsOp_Semantic<"transpose", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
let summary = "N-D transpose with arbitrary permutation (ONNX Transpose)";
let description = [{
Permutes the dimensions of the input tensor according to the `perm`
Expand Down Expand Up @@ -1831,11 +1840,12 @@ def Hip_TransposeOp : Hip_DpsOp<"transpose", /*traits=*/[],
let hasVerifier = 1;
}

def Hip_GatherOp : Hip_DpsOp<"gather", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
def Hip_GatherOp :
Hip_DpsOp_Semantic<"gather", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
let summary = "Gather elements along an axis";
let description = [{
Gathers elements from data tensor along the given axis using indices.
Expand All @@ -1862,6 +1872,7 @@ def Hip_GatherOp : Hip_DpsOp<"gather", /*traits=*/[],
`outs` `(` $output `:` type($output) `)`
attr-dict (`:` type($result_tensors)^)?
}];
let hasVerifier = 1;
}

def Hip_OneHotOp : Hip_DpsOp<"one_hot", /*traits=*/[],
Expand Down Expand Up @@ -1978,11 +1989,8 @@ def Hip_ScatterElementsOp : Hip_DpsOp<"scatter_elements", /*traits=*/[],
}];
}

def Hip_GatherElementsOp : Hip_DpsOp<"gather_elements", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
def Hip_GatherElementsOp :
Hip_DpsOp_SameShape<"gather_elements", /*sourceAccessor=*/"Indices"> {
let summary = "Gather elements along an axis using same-rank indices";
let description = [{
Gathers values from the data tensor at positions specified by the
Expand Down Expand Up @@ -4174,19 +4182,21 @@ def Hip_SliceOp :
let hasVerifier = 1;
}

def Hip_GatherNDOp : Hip_DpsOp<"gather_nd", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
def Hip_GatherNDOp :
Hip_DpsOp_Semantic<"gather_nd", /*traits=*/[],
/*outsAccessor=*/"Output",
/*autoReify=*/0,
/*autoInfer=*/1,
/*declareInfer=*/1> {
let summary = "GatherND: gather slices using N-D index tuples";
let description = [{
Gathers slices of `data` indexed by the trailing dimension of `indices`.
Implements the standard ONNX GatherND operator.

`indices` is an integer tensor whose last dimension defines the rank of
each index tuple, optionally with `batch_dims` leading batch dims that
must match `data`. Output rank is
`indices` must be an i64 tensor because the runtime and kernel ABI read
every index as `int64_t` and carry no index-width argument. Its last
dimension defines the rank of each index tuple, optionally with
`batch_dims` leading batch dims that must match `data`. Output rank is
`q + r - indices_shape[-1] - 1 - batch_dims`, where `q = rank(indices)`
and `r = rank(data)`.

Expand All @@ -4212,6 +4222,7 @@ def Hip_GatherNDOp : Hip_DpsOp<"gather_nd", /*traits=*/[],
`outs` `(` $output `:` type($output) `)`
attr-dict (`:` type($result_tensors)^)?
}];
let hasVerifier = 1;
}

def Hip_ScatterNDOp : Hip_DpsOp<"scatter_nd", /*traits=*/[],
Expand Down Expand Up @@ -4343,11 +4354,12 @@ def Hip_ModOp : Hip_DpsOp_Broadcast<"mod", /*operandGetters=*/["Lhs", "Rhs"]> {
}];
}

def Hip_SizeOp : Hip_DpsOp<"size", /*traits=*/[],
/*outsAccessor=*/"Y",
/*autoReify=*/1,
/*autoInfer=*/1,
/*declareInfer=*/1> {
def Hip_SizeOp :
Hip_DpsOp_Semantic<"size", /*traits=*/[],
/*outsAccessor=*/"Y",
/*autoReify=*/1,
/*autoInfer=*/1,
/*declareInfer=*/1> {
let summary = "Total element count of a tensor (ONNX Size, dynamic-shape path)";
let description = [{
Implements the runtime path of the standard ONNX Size operator
Expand Down Expand Up @@ -4386,6 +4398,7 @@ def Hip_SizeOp : Hip_DpsOp<"size", /*traits=*/[],
`outs` `(` $y `:` type($y) `)`
attr-dict (`:` type($result_tensors)^)?
}];
let hasVerifier = 1;
}

def Hip_NonZeroOp : Hip_DpsOp<"nonzero", /*traits=*/[], /*outsAccessor=*/"Y",
Expand Down Expand Up @@ -4527,7 +4540,9 @@ def Hip_PoolOp :
let hasVerifier = 1;
}

def Hip_ResizeOp : Hip_DpsOp<"resize"> {
def Hip_ResizeOp :
Hip_DpsOp_Semantic<"resize", /*traits=*/[],
/*outsAccessor=*/"Output", /*autoReify=*/0> {
let summary = "Spatial resize / interpolation (ONNX Resize)";
let description = [{
Implements a subset of ONNX Resize that covers the common image / volume
Expand All @@ -4548,7 +4563,9 @@ def Hip_ResizeOp : Hip_DpsOp<"resize"> {
Per-axis scale is recovered at runtime from the input/output spatial
extents (`scale = in_dim / out_dim`), so neither `scales` nor `sizes`
appear as operands here — they were resolved at compile time into the
output type that the upstream importer produced.
output type that the upstream importer produced. Therefore every output
spatial extent must be static. N/C come from the input and must agree with
the output template.

Uses destination-passing style: the output buffer is provided.

Expand All @@ -4574,6 +4591,8 @@ def Hip_ResizeOp : Hip_DpsOp<"resize"> {
`outs` `(` $output `:` type($output) `)`
attr-dict (`:` type($result_tensors)^)?
}];

let hasVerifier = 1;
}

// hip.readback_dim : materialize a kernel-computed runtime extent as a host
Expand Down
Loading
Loading