Parallelize kernel JIT compilation

Complete example: loom/binding/c/example/product_frontier.c

Loading a model or specializing a workload can make dozens or hundreds of kernels ready to compile at once. Compiling them serially extends startup latency, while giving every kernel its own compiler or thread pool repeats setup and oversubscribes the host.

The loomc API turns the live kernel set into small, independent build requests. An application checks each request against its cache, submits only the misses to a bounded scheduler, and binds the completed kernel products back to the command program that requested them. Immutable compiler state is shared across the process; each worker owns only its mutable scratch workspace.

The scheduler boundary is intentionally generic. Applications with an existing executor implement loomc_task_sink_t. Applications without one can use Loom's optional loomc_task_pool_t and loomc_task_queue_t. The compile task and its ownership rules are identical in both cases.

This guide covers:

• when independent kernel compilation improves JIT latency;
• how command construction discovers only the kernels that remain live after linking and specialization;
• where cache lookup, scheduling, compilation, and completion belong; and
• how to use either an application scheduler or Loom's standard worker pool.

Choose this path for batches and continuous specialization

Independent scheduling helps when one application event creates several compilation opportunities:

Application event	Useful scheduling policy
Model load	Compile the model's live kernel set concurrently and retain the loaded results with the model plan.
Shape or workload specialization	Reuse cached versions and schedule only newly selected semantic classes.
Continuous JIT	Keep a bounded compile queue alive and attach each invocation to its own completion object.
Autotuning or search	Submit independent configuration candidates without constructing a compiler per candidate.

A single kernel compiled once during process startup does not need this machinery. Ahead-of-time packaging remains a better fit when all relevant target and workload facts are already known at build time.

Split command construction from kernel compilation

A command program describes host-side orchestration and names the executable kernel entries it will dispatch. Building that command program can determine which kernel implementations are actually needed without compiling every kernel body in the source catalog.

command roots + configuration
             |
             v
loomc_cmd_program_product_build
             |
             +----> portable command bytes
             |
             +----> executable-entry requirements
                         |
                         +----> source-backed request -> cache -> compile
                         |
                         +----> external entry -------> application binding

The API reference calls this split a product frontier. A product is the immutable output of one successful compiler operation. Its frontier is simply the set of requirements that another product or the application must satisfy. There is no frontier object to manage and no special source format: every published kernel request contains ordinary Loom bytecode plus exact roots.

Command construction applies linking, configuration, dead-code elimination, and schedule selection before publishing requests. Dead branches create no JIT work. Launch sites that select the same semantic kernel class can share one request. A bodyless kernel.entry.decl remains an external requirement and does not create a fake compile job.

Prepare a long-lived compile service

Prepare expensive immutable state once:

• one context and target environment;
• source catalogs frozen into link indexes;
• compilers and pass programs for the supported product routes; and
• a bounded scheduler with one loomc_workspace_t per worker.

The checked example uses Loom's standard four-worker pool and attaches one compile queue:

static loomc_status_t prepare_compile_scheduler(jit_scheduler_t* scheduler) {
  loomc_status_t status = loomc_compiler_create(
      scheduler->context, NULL, scheduler->allocator, &scheduler->compiler);

  loomc_result_t* pass_result = NULL;
  if (loomc_status_is_ok(status)) {
    status = loomc_pass_program_create_from_pipeline_text(
        scheduler->context, loomc_make_cstring_view("canonicalize,cse"), NULL,
        scheduler->allocator, &scheduler->pass_program, &pass_result);
  }
  if (loomc_status_is_ok(status)) {
    status =
        require_successful_result(pass_result, "pipeline preparation failed");
  }
  loomc_result_release(pass_result);

  const loomc_task_pool_options_t pool_options = {
      .type = LOOMC_STRUCTURE_TYPE_TASK_POOL_OPTIONS,
      .structure_size = sizeof(pool_options),
      .max_worker_count = 4,
  };
  if (loomc_status_is_ok(status)) {
    status = loomc_task_pool_allocate(&pool_options, scheduler->allocator,
                                      &scheduler->task_pool);
  }
  if (loomc_status_is_ok(status)) {
    status = loomc_task_queue_allocate(
        scheduler->task_pool, scheduler->allocator, &scheduler->compile_queue);
  }
  if (loomc_status_is_ok(status)) {
    scheduler->compile_sink = loomc_task_queue_sink(scheduler->compile_queue);
    scheduler->worker_workspace_count =
        loomc_task_pool_worker_count(scheduler->task_pool);
    status = require_condition(scheduler->worker_workspace_count != 0,
                               "task pool created no workers");
  }
  if (loomc_status_is_ok(status)) {
    status = loomc_allocator_malloc(scheduler->allocator,
                                    scheduler->worker_workspace_count *
                                        sizeof(*scheduler->worker_workspaces),
                                    (void**)&scheduler->worker_workspaces);
  }
  loomc_host_size_t initialized_workspace_count = 0;
  while (loomc_status_is_ok(status) &&
         initialized_workspace_count < scheduler->worker_workspace_count) {
    status = loomc_workspace_create(
        NULL, scheduler->allocator,
        &scheduler->worker_workspaces[initialized_workspace_count]);
    if (loomc_status_is_ok(status)) ++initialized_workspace_count;
  }
  scheduler->worker_workspace_count = initialized_workspace_count;
  return status;
}

loomc_task_pool_worker_count reports the actual worker count after host CPU affinity and topology are applied. The example allocates exactly that many workspaces. Task callbacks receive a dense, mutually exclusive worker ordinal, so selecting scratch is an array lookup rather than a lock or workspace lease.

An application-owned executor replaces only the sink and worker setup. It implements loomc_task_sink_t::submit, assigns each execution lane a stable ordinal, and invokes loomc_task_execute exactly once for every accepted task. The application does not link loomc/task_pool.h, loomc/task_queue.h, or the IREE task runtime in that configuration.

Publish compile requests while building the command

Pass a loomc_request_sink_t when building the selected command roots:

static loomc_status_t build_command_product(jit_scheduler_t* scheduler,
                                            loomc_host_size_t root_ordinal) {
  const loomc_cmd_program_product_options_t options = {
      .type = LOOMC_STRUCTURE_TYPE_CMD_PROGRAM_PRODUCT_OPTIONS,
      .structure_size = sizeof(options),
      .link_index = scheduler->link_index,
      .root_symbol_ordinals = &root_ordinal,
      .root_symbol_count = 1,
      .request_sink =
          {
              .publish = schedule_kernel_request,
              .user_data = scheduler,
          },
  };
  loomc_result_t* result = NULL;
  loomc_status_t status = loomc_cmd_program_product_build(
      scheduler->command_workspace, &options, scheduler->allocator,
      &scheduler->command_product, &result);
  if (loomc_status_is_ok(status)) {
    status = require_successful_result(result, "command construction failed");
  }
  if (loomc_status_is_ok(status)) {
    status = require_condition(
        loomc_cmd_program_product_program_count(scheduler->command_product) ==
                1 &&
            loomc_product_requirement_count(scheduler->command_product) ==
                COMMAND_REQUIREMENT_CAPACITY &&
            scheduler->kernel_compile_count == KERNEL_COMPILE_CAPACITY,
        "command construction published an unexpected kernel set");
  }
  loomc_result_release(result);
  return status;
}

loomc_cmd_program_product_build is synchronous with respect to the command product, but calls request_sink.publish as independent kernel requests become available. The callback may submit work immediately, allowing kernel compilation to overlap the remainder of command construction.

The example command uses three source-backed kernels and one externally provided entry. Command construction therefore publishes three requests while returning four executable-entry requirements. Omitting request_sink.publish builds the same command bytes and requirements without opening kernel implementation bodies; that is useful when all executables are already available from another provider.

Publication is provisional until command construction returns both an OK status and a succeeded result. A latency-oriented service can schedule requests immediately and discard their parent bindings if the command later fails. A service that cannot tolerate speculative work can retain the requests in the callback and submit them only after the parent succeeds.

Check the cache before crossing the scheduler

The request callback is the cache boundary. A production key normally includes:

• request bytecode contents and exact root ordinals;
• the required product route;
• target profile and invocation configuration;
• selected pass pipeline and artifact options; and
• compiler and executable-format compatibility versions.

A cache hit records the existing product against the request's parent bindings without allocating a task. A miss wraps the immutable request in an application-owned task record and submits it. Cache policy stays outside loomc: an in-process JIT, a compiler service, and a persistent artifact cache can use different storage and eviction policies over the same requests.

The checked callback omits a cache because every process run starts empty. Its scheduler-facing path is the same path used after a real cache miss:

static loomc_status_t schedule_kernel_request(void* user_data,
                                              loomc_request_t* request) {
  jit_scheduler_t* scheduler = (jit_scheduler_t*)user_data;
  if (loomc_request_product_descriptor(request) !=
      loomc_compiled_module_product_descriptor()) {
    loomc_request_release(request);
    return loomc_make_status(LOOMC_STATUS_UNIMPLEMENTED,
                             "no compiler route for published request");
  }
  if (scheduler->kernel_compile_count == KERNEL_COMPILE_CAPACITY) {
    loomc_request_release(request);
    return loomc_make_status(LOOMC_STATUS_RESOURCE_EXHAUSTED,
                             "example kernel compile capacity exceeded");
  }

  kernel_compile_output_t* output =
      &scheduler->kernel_compiles[scheduler->kernel_compile_count];
  output->request = request;

  loomc_task_t* task = NULL;
  loomc_status_t status =
      allocate_kernel_compile_task(scheduler, output, &task);
  if (loomc_status_is_ok(status)) {
    status = loomc_task_sink_submit(scheduler->compile_sink, task);
  }
  if (!loomc_status_is_ok(status)) {
    if (task != NULL) loomc_task_destroy(task);
    loomc_request_release(output->request);
    memset(output, 0, sizeof(*output));
    return status;
  }

  // The sink owns `task` after successful submission. An inline scheduler may
  // already have executed and destroyed it, so only the output slot is used.
  ++scheduler->kernel_compile_count;
  return loomc_ok_status();
}

The publication callback owns request at entry. Successful task submission transfers the task to the sink, which executes and destroys it exactly once. A rejected submission leaves task ownership with the caller. The output slot retains the request until its completed product has been bound or discarded.

The request's process-local product descriptor selects the compilation route. This example accepts the compiled-module route used by loomc_compile_request. A service supporting additional product families uses the descriptor as a dispatch key rather than guessing from symbol names, filenames, or target strings.

Compile with worker-local scratch

Each task shares the prepared compiler and pass program and selects its workspace by worker ordinal:

static void execute_kernel_compile_task(loomc_task_t* base_task,
                                        loomc_host_size_t worker_ordinal) {
  kernel_compile_task_t* task = (kernel_compile_task_t*)base_task;
  jit_scheduler_t* scheduler = task->scheduler;
  kernel_compile_output_t* output = task->output;
  if (worker_ordinal >= scheduler->worker_workspace_count) {
    output->status = loomc_make_status(
        LOOMC_STATUS_OUT_OF_RANGE,
        "scheduler supplied an invalid compiler worker ordinal");
    return;
  }

  loomc_workspace_t* workspace = scheduler->worker_workspaces[worker_ordinal];
  const loomc_compile_options_t options = {
      .type = LOOMC_STRUCTURE_TYPE_COMPILE_OPTIONS,
      .structure_size = sizeof(options),
      .module_name = loomc_make_cstring_view("scheduled-kernel"),
      .artifact_flags = LOOMC_COMPILE_ARTIFACT_FLAG_MODULE_BYTECODE,
  };
  output->status = loomc_compile_request(
      scheduler->compiler, workspace, scheduler->pass_program, output->request,
      &options, scheduler->allocator, &output->product, &output->result);
  loomc_workspace_trim(workspace);
}

static void destroy_kernel_compile_task(loomc_task_t* base_task) {
  kernel_compile_task_t* task = (kernel_compile_task_t*)base_task;
  loomc_allocator_free(task->scheduler->allocator, task);
}

static const loomc_task_vtable_t kKernelCompileTaskVtable = {
    .execute = execute_kernel_compile_task,
    .destroy = destroy_kernel_compile_task,
};

The example requests compiled .loombc so it runs in target-independent builds. A native JIT gives the task its prepared target pipeline and continues through the compile-and-emit sequence shown in Embed kernel JIT compilation. The scheduling contract does not change when the result becomes an HSACO, SPIR-V module, or another runtime executable.

Task execution writes its loomc_status_t, loomc_result_t, and product into caller-owned completion state. A non-OK status reports API misuse or an infrastructure failure. An OK status may still accompany a failed compiler result with diagnostics. Both must be checked before caching or loading a product.

Bind completed kernels to the command

Every request carries provisional mappings from requirement ordinals in the parent command product to root ordinals in the child request. Successful child products preserve request-root order as export order, so the same mapping identifies the compiled export that satisfies each command requirement:

static loomc_status_t inspect_kernel_products(jit_scheduler_t* scheduler,
                                              bool parent_succeeded) {
  loomc_status_t status = loomc_ok_status();
  bool resolved_requirements[COMMAND_REQUIREMENT_CAPACITY] = {false};
  loomc_host_size_t resolved_requirement_count = 0;

  for (loomc_host_size_t i = 0; i < scheduler->kernel_compile_count; ++i) {
    kernel_compile_output_t* output = &scheduler->kernel_compiles[i];
    const bool compile_call_succeeded = loomc_status_is_ok(output->status);
    status = loomc_status_join(status, output->status);
    output->status = loomc_ok_status();
    if (!compile_call_succeeded) continue;

    if (output->result == NULL || !loomc_result_succeeded(output->result)) {
      print_result_diagnostics(output->result);
      status = loomc_status_join(
          status, loomc_make_status(LOOMC_STATUS_FAILED_PRECONDITION,
                                    "a scheduled kernel failed compilation"));
      continue;
    }

    const loomc_artifact_t* artifact =
        loomc_product_artifact_at(output->product, 0);
    const bool product_is_valid =
        loomc_product_descriptor(output->product) ==
            loomc_compiled_module_product_descriptor() &&
        loomc_product_export_count(output->product) ==
            loomc_request_root_count(output->request) &&
        loomc_product_requirement_count(output->product) == 0 &&
        loomc_product_artifact_count(output->product) == 1 &&
        artifact != NULL &&
        loomc_string_view_equal(
            artifact->format,
            loomc_make_cstring_view(LOOMC_ARTIFACT_FORMAT_LOOM_BYTECODE));
    if (!product_is_valid) {
      status = loomc_status_join(
          status,
          loomc_make_status(LOOMC_STATUS_FAILED_PRECONDITION,
                            "a scheduled kernel returned a bad product"));
      continue;
    }

    if (!parent_succeeded) continue;
    const loomc_host_size_t binding_count =
        loomc_request_binding_count(output->request);
    for (loomc_host_size_t j = 0; j < binding_count; ++j) {
      loomc_request_binding_t binding = {0};
      const bool binding_is_valid =
          loomc_request_binding_at(output->request, j, &binding) &&
          binding.requirement_ordinal < COMMAND_REQUIREMENT_CAPACITY &&
          binding.root_ordinal < loomc_product_export_count(output->product) &&
          !resolved_requirements[binding.requirement_ordinal];
      if (!binding_is_valid) {
        status = loomc_status_join(
            status,
            loomc_make_status(LOOMC_STATUS_FAILED_PRECONDITION,
                              "a kernel request returned an invalid binding"));
        continue;
      }
      resolved_requirements[binding.requirement_ordinal] = true;
      ++resolved_requirement_count;
    }
  }

  if (!parent_succeeded || !loomc_status_is_ok(status)) return status;

  loomc_host_size_t external_requirement_count = 0;
  for (loomc_host_size_t i = 0; i < COMMAND_REQUIREMENT_CAPACITY; ++i) {
    if (resolved_requirements[i]) continue;
    loomc_cmd_entry_requirement_t requirement = {0};
    if (!loomc_cmd_program_product_entry_requirement_at(
            scheduler->command_product, i, &requirement) ||
        !loomc_string_view_equal(requirement.symbol,
                                 loomc_make_cstring_view("external"))) {
      return loomc_make_status(LOOMC_STATUS_FAILED_PRECONDITION,
                               "unexpected external kernel requirement");
    }
    ++external_requirement_count;
  }

  loomc_cmd_program_t program = {0};
  if (!loomc_cmd_program_product_program_at(scheduler->command_product, 0,
                                            &program) ||
      resolved_requirement_count != KERNEL_COMPILE_CAPACITY ||
      external_requirement_count != 1) {
    return loomc_make_status(LOOMC_STATUS_FAILED_PRECONDITION,
                             "kernel products do not satisfy the command");
  }
  printf("command=%.*s requirements=%zu compiled=%zu external=%zu\n",
         (int)program.symbol.size, program.symbol.data,
         (size_t)loomc_product_requirement_count(scheduler->command_product),
         (size_t)resolved_requirement_count,
         (size_t)external_requirement_count);
  return status;
}

The example verifies three compiled bindings and leaves the bodyless @external entry for the application. A real runtime uses those mappings to join loaded executable exports to command-program entry slots. Two command products may reuse one cached child product while carrying different parent requirement ordinals.

Completion belongs to the application request, not to the worker pool. The short-lived example shuts down its queue to drain one finite batch. A continuous JIT keeps the pool and queues alive, increments a per-command pending count for each accepted request, resolves cache hits immediately, and signals its future, callback, or event-loop message when every required binding has completed. Observing an empty queue is not a completion protocol because another producer may submit work at any time.

Run the checked example

python dev.py bazel run //loom/binding/c/example:product_frontier

The output distinguishes compiled source kernels from the executable that the application must provide:

command=run requirements=4 compiled=3 external=1

For the generic task protocol and standard pool lifecycle in isolation, see Use the JIT task pool. For target profiles, native artifact emission, launch configuration, and runtime loading, continue with Embed kernel JIT compilation. The generated loomc C API reference defines the exact ownership and thread-safety contract for every handle used here.