Skip to content

HloGenerator.generate emits an empty module: it doesn't pushTape onto Execution.tapeStack around forward #668

Description

@michalharakal

HloGenerator.generate emits an empty module — it doesn't push the tape onto Execution.tapeStack

HloGenerator.generateBlocking(model, input, name) returns a structure-only module
(func.func @<name>() -> () { return }, no stablehlo.* ops) for any Model whose
forward records through the tensor-bound ops — because it never activates the tape on
Execution.tapeStack, so the recorded operations don't land on the tape and
toComputeGraph() is empty. (This is the root cause of the "structure-only export"
noted in #663.)

The working recipe (from SKaiNET-transformers GemmaMlirDumpTest)

The transformers repo traces real models correctly by pushing the current tape:

val ctx = DefaultGraphExecutionContext.tape(baseOps = VoidTensorOps())
val tape = ctx.record {
    val ct = (this as DefaultGraphExecutionContext).currentTape ?: error("no tape")
    Execution.tapeStack.pushTape(ct)                 // <-- HloGenerator is missing this
    try { model.forward(input, this as ExecutionContext) }
    finally { Execution.tapeStack.popTape() }
}.first
val graph = (tape as DefaultExecutionTape).toComputeGraph(synthesizeExternalInputs = true)
val mlir = toStableHlo(graph, "gemma").content       // real ops (rmsNorm/sdpa/rope/...)

With pushTape, the same Model that HloGenerator exports as empty produces a full graph.

Ask

HloGenerator.generate (and generateBlocking) should push/pop the recording tape on
Execution.tapeStack around the forward call (and use
toComputeGraph(synthesizeExternalInputs = true)), so the documented one-call API matches
the working manual recipe. A test that asserts a non-empty module for Rgb2GrayScaleMatMul
(or any tensor-ops model) would lock it in.

cc #663

Activity

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

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions