Internal error during proc inlining

[rw1nkler]

Describe the bug

Inlining a proc that spawns another proc and contains a logic that performs operations on tokens, causes an internal error.

To Reproduce

One can use the following DSLX code to observe the issue:

// DSLX code

proc Passthrough {
    data_r: chan<u32> in;
    data_s: chan<u32> out;

    config(data_r: chan<u32> in, data_s: chan<u32> out) { (data_r, data_s) }

    init { () }

    next(state: ()) {
        let (tok, data) = recv(join(), data_r);
        let tok = send(tok, data_s, data);
    }
}

proc PassthroughNested {
    intr_r: chan<u32> in;
    data_s: chan<u32> out;
    config(data_r: chan<u32> in, data_s: chan<u32> out) {
        let (intr_s, intr_r) = chan<u32, u32:1>("intr");
        spawn Passthrough(data_r, intr_s);
        (intr_r, data_s)
    }

    init { () }
    next(tok: token, state: ()) {
        let (tok, data) = recv(tok, intr_r);
        let tok = send(tok, data_s, data);
    }
}

Here is a bazel rule that can be used to observe the problem:

# BUILD rules

xls_dslx_opt_ir(
    name = "passthrough_nested_opt_ir",
    library = "passthrough_dslx",
    dslx_top = "PassthroughNested",
    opt_ir_args = { "inline_procs": "true" },
)

Building the passthrough_nested_opt_ir target causes the following error:

ERROR: /home/rwinkler/xls/xls/examples/BUILD:1031:16: Optimizing IR file: bazel-out/k8-fastbuild/bin/xls/examples/passthrough_nested_opt_ir.ir failed: (Exit 1): opt_main failed: error executing command (from target //xls/examples:passthrough_nested_opt_ir) bazel-out/k8-opt-exec-2B5CBBC6/bin/xls/tools/opt_main bazel-out/k8-fastbuild/bin/xls/examples/passthrough_nested_opt_ir.ir --inline_procs true --output_path ... (remaining 3 arguments skipped)

Use --sandbox_debug to see verbose messages from the sandbox and retain the sandbox build root for debugging
E0604 13:45:56.167111   10260 proc_inlining_pass.cc:1388] INTERNAL: XLS_RET_CHECK failure (xls/passes/proc_inlining_pass.cc:1388) !activations_in.empty()
0x559d6610768a: xabsl::StatusBuilder::CreateStatusAndConditionallyLog()
0x559d65f230c9: absl::lts_20240116::StatusOr<>::StatusOr<>()
0x559d65f225e1: xls::(anonymous namespace)::ProcThread::AllocateActivationNode()
0x559d65f1e731: xls::(anonymous namespace)::ProcThread::Create()
0x559d65f0ccca: xls::(anonymous namespace)::InlineProcAsProcThread()
0x559d65f079bd: xls::ProcInliningPass::RunInternal()
0x559d65db9889: xls::PassBase<>::Run()
0x559d65dba882: xls::CompoundPassBase<>::RunNested()
0x559d65dba724: xls::CompoundPassBase<>::RunNested()
0x559d65dba0ae: xls::CompoundPassBase<>::RunInternal()
0x559d65db9889: xls::PassBase<>::Run()
0x559d65db48ac: xls::tools::OptimizeIrForTop()
0x559d65db6677: xls::tools::OptimizeIrForTop()
0x559d65d6d629: main
0x7fc2481ebd90: [unknown]

Error: INTERNAL: XLS_RET_CHECK failure (xls/passes/proc_inlining_pass.cc:1388) !activations_in.empty() ; Running pass #208: Proc inlining [short: proc_inlining]; Running pass #208: Post-inlining passes [short: post-inlining]; Running pass #208: Top level pass pipeline [short: ir]
Target //xls/examples:passthrough_nested_opt_ir failed to build

What is interesting, a procs with empty next(), like the one below, can be inlined without any problems:

proc PassthroughWrapped {
    config(data_r: chan<u32> in, data_s: chan<u32> out) {
        let (intr_s, intr_r) = chan<u32, u32:1>("intr");
        spawn Passthrough(data_r, intr_s);
        spawn Passthrough(intr_r, data_s);
    }

    init { () }
    next(state: ()) { }
}

Expected behavior

It should be possible to inline the procs