Tutorial: Finding the SDK Kernel for an Op CGC Cannot Lower

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/custom_op/plugin_tutorial/plugin_autofinder.ipynb.

Finding the SDK kernel for an op CGC cannot lower

The SDK's sweep tests under sdk/tests/sweep_tests/<kernel>/ pair a CCL .cpp that calls an SDK kernel with the graph that kernel stands in for. The autofinder ships a reference built from those tests inside sdk_cli (sdk_cli/cgc_plugin/autofinder/sdk_sweep_reference.json), so no SDK checkout is opened here. This notebook is the dry run of that tool on the voxel pooling model from the voxel pooling tutorial, whose VoxelPooling node has no ONNX schema and no convert-map entry: it finds the SDK kernel for the node, generates a bundle whose adapter forwards to that kernel and substitutes the node. It stops before compiling. The compile, the ISS run and the check against the SDK's own NumPy reference are the second half of the voxel pooling tutorial, which runs this flow end to end on the same model.

1. Setup

Imports the autofinder modules. load_reference() reads sdk_sweep_reference.json from the package. staleness_warning compares the SDK commit the reference was built from with the toolchain in use ($SDK_HASH on the CI image, the checkout behind the installed quadric package here) and returns a message when they differ. Creates _autofinder_out/ fresh and prints the Python, the plugin package path, the reference description, the toolchain check and the model path.

import os
import shutil
import sys
from collections import Counter

import onnx

from sdk_cli import cgc_plugin
from sdk_cli.cgc_plugin.autofinder.map_sdk_tests import (
    candidate_op_type,
    graph_grounded_candidates,
)
from sdk_cli.cgc_plugin.autofinder.propose_bundles import propose
from sdk_cli.cgc_plugin.autofinder.reference import (
    describe_reference,
    load_reference,
    staleness_warning,
)
from sdk_cli.utils.model_helpers import example_path

tutorial_dir = example_path("custom_op/plugin_tutorial")
model_path = example_path("custom_op/voxel_pooling/voxel_original.onnx")

## Everything this notebook writes lands here, rebuilt each run.
work_dir = tutorial_dir / "_autofinder_out"
shutil.rmtree(work_dir, ignore_errors=True)
work_dir.mkdir()

## The reference ships inside sdk_cli; no SDK checkout is opened.
reference = load_reference()
stale = staleness_warning(reference)
print(f"python     : {sys.executable}")
print(f"cgc_plugin : {cgc_plugin.__file__}")
print(f"reference  : {describe_reference(reference)}")
print(f"toolchain  : {stale or 'built from the same SDK the reference describes'}")
print(f"model      : {model_path}")

python     : /usr/bin/python3
cgc_plugin : /quadric/sdk-cli/sdk_cli/cgc_plugin/__init__.py
reference  : sdk@3f1b472d0 on devel (detached), 275 sweep-test dirs
toolchain  : the shipped sweep reference describes sdk@3f1b472d0 but the toolchain in use is sdk@8b91cc4; kernels may have moved. Rebuild it: python -m sdk_cli.cgc_plugin.autofinder.build_reference
model      : /quadric/sdk-cli/examples/custom_op/voxel_pooling/voxel_original.onnx

2. Which op-types CGC cannot lower

cgc_native_ops() returns the op-types in CGC's convert map. discover_candidate_ops(model, native_ops) groups the model's nodes by op-type and keeps the types not in that set. For each, pattern_from_nodes prints the SSA topology pattern written from the first node. The assertion pins VoxelPooling as the only such op-type.

from sdk_cli.cgc_plugin import cgc_native_ops, discover_candidate_ops, pattern_from_nodes

model = onnx.load(str(model_path))
native_ops = cgc_native_ops()
candidates = discover_candidate_ops(model, native_ops=native_ops)

print(
    f"{model_path.name}: {len(model.graph.node)} nodes, op-types {sorted({n.op_type for n in model.graph.node})}"
)
print(f"convert map         : {len(native_ops)} op-types CGC lowers natively")
print(f"NOT natively lowered: { {op: len(nodes) for op, nodes in candidates.items()} }")
for op_type, nodes in candidates.items():
    pattern, output = pattern_from_nodes(model, nodes[:1])
    print(f"\npattern for {op_type}, by example from node '{nodes[0].name}' (output = {output}):")
    print(pattern)

assert list(candidates) == ["VoxelPooling"], candidates

voxel_original.onnx: 7 nodes, op-types ['Cast', 'Constant', 'DequantizeLinear', 'Identity', 'QuantizeLinear', 'VoxelPooling']
convert map         : 169 op-types CGC lowers natively
NOT natively lowered: {'VoxelPooling': 1}

pattern for VoxelPooling, by example from node '/VoxelPooling' (output = voxelpooling0):
in0 = *
in1 = *
voxelpooling0 = VoxelPooling(in0, in1)

3. Compile without a plugin

Builds the plugin's ChimeraJob with plugins=False and compiles in a scratch folder (CGC writes ccl_build/ into the working directory). The compile fails; the exception is kept and the CGCUnsupported.yaml report CGC wrote is printed. The assertions require the failure and that the report names VoxelPooling.

from sdk_cli.cgc_plugin.chimera_job import ChimeraJob  # the plugin's ChimeraJob, not tvm's

check_dir = work_dir / "no_plugin"
check_dir.mkdir()
previous_cwd = os.getcwd()
os.chdir(check_dir)  # CGC writes ccl_build/ into the cwd
compile_error = None
try:
    ChimeraJob(str(model_path), plugins=False).compile(quiet=True)
except Exception as exc:  # noqa: BLE001 - the failure IS the measurement
    compile_error = exc
finally:
    os.chdir(previous_cwd)

assert compile_error is not None, "CGC must reject VoxelPooling without a plugin"
unsupported_report = next(check_dir.glob("ccl_build/*/compile_errors/CGCUnsupported.yaml"))
print(f"compile WITHOUT plugin: {type(compile_error).__name__}: {compile_error}")
print(f"\n{unsupported_report.relative_to(check_dir)}:")
print(unsupported_report.read_text())
assert "VoxelPooling" in unsupported_report.read_text()

/usr/local/lib/python3.10/dist-packages/tvm/relay/frontend/onnx.py:6495: UserWarning: No Op registered for VoxelPooling with domain_version of 16

==> Context: Bad node spec for node. Name: /VoxelPooling OpType: VoxelPooling
  warnings.warn(str(e))
/usr/local/lib/python3.10/dist-packages/tvm/relay/frontend/onnx.py:6201: UserWarning: Operator VoxelPooling not implemented.
  warnings.warn(f"Operator {op_name} not implemented.")
WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: output
WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: scale
WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: zero_point
/usr/local/lib/python3.10/dist-packages/tvm/relay/frontend/onnx.py:6495: UserWarning: No Op registered for VoxelPooling with domain_version of 16

==> Context: Bad node spec for node. Name: CustomOp/customOp00 OpType: QuadricCustomOp
  warnings.warn(str(e))


[SDK-CLI] : Unsupported nodes:
['/VoxelPooling']
Check the file: .//quadric/sdk-cli/examples/custom_op/plugin_tutorial/_autofinder_out/no_plugin/ccl_build/voxel_original_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/compile_errors/CGCUnsupported.yaml for more detail.
compile WITHOUT plugin: Exception: check above file for more information

ccl_build/voxel_original_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/compile_errors/CGCUnsupported.yaml:
release: $(git rev-parse HEAD)
model: /quadric/sdk-cli/examples/custom_op/voxel_pooling/voxel_original.onnx
product: QC-U
ocm_size: 16MB
error_list:
- filter: VoxelPooling
  code: QE99
  node_name: /VoxelPooling

4. The sweep-test reference

Reads the kernels table of the reference. Counts the sweep directories, those calling an SDK symbol, the distinct symbols, and splits the symbols by whether candidate_op_type(symbol) is in the convert map. graph_grounded_candidates lists the directories whose paired .onnx holds an op-type the convert map lacks. Prints the directory count with a paired graph and the count whose kernel call the adapter parsed. The last lines look up voxelPooling: not in the convert map, and shipped by the voxel_pooling sweep directory.

kernels = reference["kernels"]
symbols = sorted({symbol for info in kernels.values() for symbol in info["symbols"]})
covered = [s for s in symbols if candidate_op_type(s) in native_ops]
uncovered = [s for s in symbols if candidate_op_type(s) not in native_ops]
grounded = graph_grounded_candidates(kernels, set(native_ops))

print(f"sweep-test dirs in the reference      : {len(kernels)}")
print(
    f"  calling an SDK kernel symbol        : {sum(1 for info in kernels.values() if info['symbols'])}"
)
print(f"distinct SDK symbols                  : {len(symbols)}")
print(f"  op-type IS in the convert map       : {len(covered)}  (CGC lowers it; no bundle)")
print(
    f"  op-type NOT in the convert map      : {len(uncovered)}  (candidates, by symbol NAME - over-generates)"
)
print(
    f"dirs pairing a kernel with a .onnx    : {sum(1 for info in kernels.values() if info['graphs'])}"
)
print(
    f"dirs whose call the adapter parsed    : {sum(1 for info in kernels.values() if 'call' in info)}"
)
print("\nGRAPH-GROUNDED candidates - the paired .onnx holds an op the convert map lacks:")
for row in grounded:
    print(
        f"  {row['kernel_dir']:28} {', '.join(row['unmapped_op_types']):14} {', '.join(row['symbols'][:3])}"
    )

voxel_symbol = "voxelPooling"
print(
    f"\n'{voxel_symbol}' -> candidate op-type {candidate_op_type(voxel_symbol)!r}, "
    f"in convert map: {candidate_op_type(voxel_symbol) in native_ops}; "
    f"sweep dir(s): {[name for name, info in kernels.items() if voxel_symbol in info['symbols']]}"
)
assert voxel_symbol in uncovered

sweep-test dirs in the reference      : 275
  calling an SDK kernel symbol        : 160
distinct SDK symbols                  : 168
  op-type IS in the convert map       : 22  (CGC lowers it; no bundle)
  op-type NOT in the convert map      : 146  (candidates, by symbol NAME - over-generates)
dirs pairing a kernel with a .onnx    : 34
dirs whose call the adapter parsed    : 59

GRAPH-GROUNDED candidates - the paired .onnx holds an op the convert map lacks:

'voxelPooling' -> candidate op-type 'VoxelPooling', in convert map: False; sweep dir(s): ['voxel_pooling']

5. Propose a bundle

propose(model_path, out_dir, None, False) runs a dry run: for each op-type CGC cannot lower it prints the pattern, the proposed bundle name and whether the reference holds a sweep test for it. The assertions pin one proposal, VoxelPooling, found in voxel_pooling. SweepCall.from_dict rebuilds the parsed kernel call and the cell prints its symbol, its arguments, and for each tensor argument its role, dtype and shape as the sweep generator declares them.

from sdk_cli.cgc_plugin.autofinder.adapter import SweepCall

proposals = propose(model_path, work_dir / "proposed", None, False)

assert [p["op_type"] for p in proposals] == ["VoxelPooling"], proposals
proposal = proposals[0]
assert proposal["sdk_kernel"] is not None, "the sweep tests ship nn::voxelPooling"
assert proposal["sdk_kernel"]["dir"] == "voxel_pooling", proposal["sdk_kernel"]

## The kernel call the reference parsed from that sweep test.
call = SweepCall.from_dict(proposal["sdk_kernel"]["call"])
print(f"\nsweep_tests/{call.sweep_dir}: {call.symbol}({', '.join(call.args)})")
for arg in call.args:
    if arg == call.allocator:
        print(f"  {arg:12} OCM allocator")
        continue
    dtype, dims = call.meta[call.roles[arg]]
    print(f"  {arg:12} {call.roles[arg]:14} {dtype:24} {list(dims)}")

model            : /quadric/sdk-cli/examples/custom_op/voxel_pooling/voxel_original.onnx
sdk reference    : sdk@3f1b472d0 on devel (detached), 275 sweep-test dirs


WARNING:cgc_plugin:the shipped sweep reference describes sdk@3f1b472d0 but the toolchain in use is sdk@8b91cc4; kernels may have moved. Rebuild it: python -m sdk_cli.cgc_plugin.autofinder.build_reference


!! the shipped sweep reference describes sdk@3f1b472d0 but the toolchain in use is sdk@8b91cc4; kernels may have moved. Rebuild it: python -m sdk_cli.cgc_plugin.autofinder.build_reference
sweep tests known: 275 kernel dirs

op-types CGC cannot lower: {'VoxelPooling': 1}

==============================================================================
VoxelPooling  x1  ->  bundle 'voxelpooling'
==============================================================================
pattern (by example, from the first occurrence):
    in0 = *
    in1 = *
    voxelpooling0 = VoxelPooling(in0, in1)

SDK kernel : EXISTS -- sweep_tests/voxel_pooling
             symbols in it: voxelPooling
             -> FORWARD to it. That test shows the call, the tensor types it
                expects and the staging. Do not reimplement it.

(dry run -- pass --write to scaffold this bundle)

sweep_tests/voxel_pooling: nn::voxelPooling(ddrGeom, ddrFeatures, ddrOut, ocmMem)
  ddrGeom      GeomShape      DataType(EPU_TYPE.INT16) [1, 1, None, 3]
  ddrFeatures  FeaturesShape  test_case.in_type        [1, 1, None, None]
  ddrOut       OutShape       test_case.out_type       [1, None, None, None]
  ocmMem       OCM allocator

6. Generate the bundle

generate_bundle(model, nodes, call, plugins_dir, name=..., options=..., description=...) assigns the node's inputs to the kernel's arguments by rank and literal dims (the geometry's trailing 3 fixes the order here, so one C++ overload is written), renders the forwarding adapter, and adds an eliminate_identity post-pass because an Identity feeds the node and CGC's fx canonicalizer rejects an integer copy. options carries the two values the sweep test cannot know: allow_io_in_l2_mem and a 1 MB reserved_ocm rule clamped to a quarter of the part's OCM. Prints the generated kernel.hpp and plugin.yaml.

from sdk_cli.cgc_plugin.autofinder.adapter import generate_bundle

plugins_dir = work_dir / "plugins"
bundle = generate_bundle(
    model,
    candidates["VoxelPooling"],
    call,
    plugins_dir,
    name="voxel_pooling",
    options={
        "allow_io_in_l2_mem": True,
        "reserved_ocm": {"builtin": "fixed", "params": {"bytes": 1048576, "clamp": "quarter"}},
    },
    description=(
        "VoxelPooling -> quadric_plugin::voxelPoolingStub -> nn::voxelPooling, the kernel the SDK "
        "sweep test voxel_pooling/ exercises. Adapter generated by plugin_autofinder.ipynb."
    ),
)
print(f"bundle: {bundle}  files: {sorted(p.name for p in bundle.iterdir())}\n")
print("--- kernel.hpp (generated) ---")
print((bundle / "kernel.hpp").read_text())
print("--- plugin.yaml ---")
print((bundle / "plugin.yaml").read_text())

bundle: /quadric/sdk-cli/examples/custom_op/plugin_tutorial/_autofinder_out/plugins/voxel_pooling  files: ['kernel.hpp', 'plugin.yaml']

--- kernel.hpp (generated) ---
// VoxelPooling: adapter over nn::voxelPooling, generated from the SDK sweep test.
// Forwards only. One overload per operand order the sweep shapes allow; C++ picks the
// one CGC binds.
using namespace chimera;

namespace quadric_plugin {

template <std::uint32_t d0_2, typename T0, std::uint32_t d1_2, std::uint32_t d1_3, typename T1, typename OutTensor, typename OcmAllocatorType>
INLINE void voxelPoolingStub(DdrTensor<T0, 1, 1, d0_2, 3>& in0, DdrTensor<T1, 1, 1, d1_2, d1_3>& in1, OutTensor& out, OcmAllocatorType& ocmMem) {
  nn::voxelPooling(in0, in1, out, ocmMem);
}

}  // namespace quadric_plugin

--- plugin.yaml ---
name: voxel_pooling
description: VoxelPooling -> quadric_plugin::voxelPoolingStub -> nn::voxelPooling, the kernel the SDK
  sweep test voxel_pooling/ exercises. Adapter generated by plugin_autofinder.ipynb.
impl_type: ccl
kernel:
  file: kernel.hpp
  ccl_func_name: quadric_plugin::voxelPoolingStub
  element_wise: false
match:
  mode: topology
  output: voxelpooling0
  topology: |
    in0 = *
    in1 = *
    voxelpooling0 = VoxelPooling(in0, in1)
options:
  allow_io_in_l2_mem: true
  reserved_ocm:
    builtin: fixed
    params:
      bytes: 1048576
      clamp: quarter
postprocess:
  reason: Identity ['geom_xyz_identity'] feeds VoxelPooling; CGC's fx canonicalizer cannot range-propagate
    an int copy
  onnxoptimizer_passes:
  - eliminate_identity

7. Substitute

Builds the plugin's ChimeraJob with plugins_dir set to the generated folder and reads plugin_result. Loads the substituted model and prints the op-type counts before and after: VoxelPooling and Identity are gone and one QuadricCustomOp node is present. Prints that node's inputs, outputs and ccl_func_name, and the paths of the substituted model and combined header. The assertions pin the matched bundle and the counts.

substitution_job = ChimeraJob(
    str(model_path), plugins_dir=str(plugins_dir), plugin_out_dir=str(work_dir / "substituted")
)
analysis = substitution_job.plugin_result
print(analysis)

substituted = onnx.load(analysis.model_path)
before = Counter(node.op_type for node in model.graph.node)
after = Counter(node.op_type for node in substituted.graph.node)
print(f"\nop-types before : {dict(before)}")
print(f"op-types after  : {dict(after)}")

custom_ops = [node for node in substituted.graph.node if node.op_type.startswith("QuadricCustomOp")]
for node in custom_ops:
    bound_to = next(
        onnx.helper.get_attribute_value(a) for a in node.attribute if a.name == "ccl_func_name"
    )
    print(f"\n{node.op_type} '{node.name}': {list(node.input)} -> {list(node.output)}")
    print(f"  ccl_func_name = {bound_to.decode()}")
print(f"\ncompile this model : {analysis.model_path}")
print(f"with this header   : {analysis.header_path}")

assert analysis.matched_plugins == ("voxel_pooling",), analysis.matched_plugins
assert after["VoxelPooling"] == 0 and after["Identity"] == 0 and len(custom_ops) == 1, dict(after)

WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: output
WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: scale
WARNING:root:ONNX node does not have a name, deriving temp name from its outputs. Node was named to: zero_point


cgc plugin: 1 customOp(s) from voxel_pooling -> /quadric/sdk-cli/examples/custom_op/plugin_tutorial/_autofinder_out/substituted

op-types before : {'Constant': 2, 'QuantizeLinear': 1, 'Identity': 1, 'VoxelPooling': 1, 'DequantizeLinear': 1, 'Cast': 1}
op-types after  : {'Constant': 2, 'QuantizeLinear': 1, 'QuadricCustomOp': 1, 'DequantizeLinear': 1, 'Cast': 1}

QuadricCustomOp 'CustomOp/quadric_plugin::voxelPoolingStub0': ['geom_xyz', 'img_feat_with_depth_quantized'] -> ['voxel_output']
  ccl_func_name = quadric_plugin::voxelPoolingStub

compile this model : /quadric/sdk-cli/examples/custom_op/plugin_tutorial/_autofinder_out/substituted/voxel_original.onnx
with this header   : /quadric/sdk-cli/examples/custom_op/plugin_tutorial/_autofinder_out/substituted/custom_ops.hpp

Next: compile and run

_autofinder_out/plugins/voxel_pooling/ is a complete bundle: ChimeraJob(model, plugins_dir=...) compiles the model with it. The voxel pooling tutorial does that on this model, runs it on the ISS and checks the output against the SDK sweep test's NumPy reference, bit-exact on every voxel whose sum fits the graph's int8 output.