Kernel Selection for IR Graphs

Adapter from TorchLean's op-tagged IR to backend capsules.

This is intentionally a planning layer, not an evaluator. The graph remains the semantic object; the planner selects which backend contract should implement each node or node family.

def NN.Backend.IR.op? : IR.OpKind → Option BackendOp

Backend operation tag for an IR op, or none for a node that needs no backend kernel.

These tags are the same vocabulary used by backend capsules and runtime reports. If an op maps to a tag that has no capsule for the selected profile, planning fails at the backend boundary instead of silently using a broader bucket. Pointwise operators map only to a backend operation with the same meaning: negation, true division and the tanh GELU. planNode? rejects the others.

structure NN.Backend.IR.PlannedNodeKernel : Type

Backend choice for one concrete IR node.

• nodeId : ℕ
• kind : IR.OpKind
• op : BackendOp
• capsule : KernelCapsule

def NN.Backend.IR.instReprPlannedNodeKernel.repr : PlannedNodeKernel → ℕ → Std.Format

@[instance_reducible]

instance NN.Backend.IR.instReprPlannedNodeKernel : Repr PlannedNodeKernel

structure NN.Backend.IR.GraphKernelPlan : Type

Selected kernel capsules with the source IR node identity preserved.

• kernels : Array PlannedNodeKernel

def NN.Backend.IR.instReprGraphKernelPlan.repr : GraphKernelPlan → ℕ → Std.Format

@[instance_reducible]

instance NN.Backend.IR.instReprGraphKernelPlan : Repr GraphKernelPlan

def NN.Backend.IR.GraphKernelPlan.nodeIds (p : GraphKernelPlan) : Array ℕ

Node ids covered by backend kernels, in graph order.

def NN.Backend.IR.GraphKernelPlan.capsuleNames (p : GraphKernelPlan) : Array String

Selected capsule names, in graph order.

def NN.Backend.IR.planNode? (policy : KernelPolicy) (available : Array KernelCapsule) (n : IR.Node) : Except String (Option PlannedNodeKernel)

Plan a single IR node when it corresponds to runtime work.

available is the capsule registry already restricted by Availability.filterCapsules, so a graph planner filters once rather than once per node.

def NN.Backend.IR.planGraph (policy : KernelPolicy) (availability : Availability) (registry : Array KernelCapsule) (g : IR.Graph) : Except String GraphKernelPlan

Plan every runtime-relevant node in graph order.

def NN.Backend.IR.checkedPlanGraph (policy : KernelPolicy) (availability : Availability) (registry : Array KernelCapsule) (g : IR.Graph) : Except String GraphKernelPlan

Check graph well-formedness, then plan every runtime-relevant node.