Repository navigation
Trace-mode conv2d/maxPool2d/upsample2d aren't recognized by the StableHLO exporter #1219
Description
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, 3maxPool2d, 2upsample2d.
Each emits:— the allowlist has no convolution, pooling, or upsampling entry at all.// Unsupported op 'conv2d' (type=trace) for node n249_conv2d. Known names: [Embedding, SiLU, abs, add, addScalar, ..., slice, softmax, split, sqrt, ..., transpose, unsqueeze, ...]
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{}:
conv2d—mnist-cnn,yolomaxPool2d—mnist-cnnupsample2d—ops(Nearest + Bilinear, landed 0.33.0, Add upsample2d Bilinear (eager+autodiff) and traceable StableHLO lowering (Nearest + Bilinear) #771; nearest max abs err 0.0,
bilinear 1.19e-7)
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
Merge pull request #1257 from SKaiNET-developers/test/1219-trace-conv…