Skip to content

Trace-mode conv2d/maxPool2d/upsample2d aren't recognized by the StableHLO exporter #1219

Closed

Description

@michalharakal

Trace-mode conv2d/maxPool2d/upsample2d aren't recognized by the StableHLO exporter

Summary

Exporting a network{}/sequential{}/stage{}-authored model that uses the conv2d,
maxPool2d, or upsample2d layer builders, traced to a ComputeGraph and passed to
toStableHlo, silently produces an incomplete module: every node of those three types
emits an // Unsupported op '<name>' comment instead of a real stablehlo.* op, and the
export still reports success (no exception, no non-zero exit).

The same three operations already have working StableHLO lowerings via the low-level
dag{} DSL
— proven by three passing conformance models in skainet-iree-conformance
(mnist-cnn, yolo, ops). This looks like a dispatch/registration gap specific to how
the trace path records these three op names, not a missing lowering capability overall.

Repro

skainet-iree-conformance commit 6f0ecf5 (conformance/models/yolo8graph/) — an
unmodified copy of skainet-model-yolo's real Yolo8Graph/Yolo8/YoloConfig/
YoloPost, traced and exported for the first time. Run:

./gradlew :models:yolo8graph:jvmRun --args="--out build/yolo8graph.mlir"

Result:

  • Traced graph: 451 nodes, exactly the expected 6 outputs (3 heads × [64 DFL-regression
    channels, 80 class-score channels] at strides 8/16/32 for a 640×640 input) — the
    Kotlin-level tracing and shape inference is correct end to end.
  • Of those 451 nodes, 105 fail to lower: 100 conv2d, 3 maxPool2d, 2 upsample2d.
    Each emits:
    // Unsupported op 'conv2d' (type=trace) for node n249_conv2d. Known names: [Embedding,
    SiLU, abs, add, addScalar, ..., slice, softmax, split, sqrt, ..., transpose,
    unsqueeze, ...]
    
    — the allowlist has no convolution, pooling, or upsampling entry at all.

Why this looks like a dispatch gap, not a missing lowering

This repo's own conformance suite already proves all three ops work via dag{}:

So the StableHLO lowering logic for these three concepts exists and is tested — the trace
path just isn't reaching it.

Two candidate fixes (no recommendation — I don't know the internals well enough to pick)

A — register the trace-emitted names in the existing dispatch table. If dag{}'s
builders and the trace path's node shape for these ops actually match (same attributes,
operand order, weight layout), this could be a narrow addition: route the trace-emitted
"conv2d"/"maxPool2d"/"upsample2d" (type=trace) names to the same lowering code
dag{} already uses.

B — unify the two op-recording paths. If the trace-based DSL and dag{} record
genuinely different internal representations for the same operations, the more durable
fix is having the higher-level DSL's layer builders construct dag{}-level primitives
directly, so a future op gets both paths for free instead of needing two independent
implementations that can silently diverge.

Impact

Blocks any real, production-shaped CNN authored via the network{}/sequential{}/
stage{} DSL from exporting to StableHLO at all — this surfaced while trying to export
skainet-model-yolo's Yolo8Graph (a real YOLOv8 backbone+neck+3-heads model), but isn't
specific to that model.

Note on scope

This only checked export/compile-readiness (a weight-free, void-traced structural check) —
not lowering correctness. Once this exports cleanly, an oracle-backed numerical
validation (the gemma-decoder conformance model's pattern) should follow before any real
model relies on this path. This session's investigation also only found what one real
model happens to use — a broader audit comparing the trace-path dispatch table against
dag{}'s might surface more gaps of the same kind.

Activity

  1. added a commit that references this issue on Sep 2, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions