ml_switcheroo.compiler.backends.sass.synthesizer

SASS Synthesizer and Register Allocator.

This module provides the “Middle-End” logic for the SASS compiler pipeline. It bridges the gap between high-level Abstract Logic (LogicalGraph) and low-level Physical Assembly (Instruction nodes/Registers).

It contains: 1. RegisterAllocator: Map Symbolic Variables (e.g., ‘x’, ‘bias’) to

Physical Registers (e.g., ‘R0’, ‘R1’).

  1. SassSynthesizer: - Target Transformation (`from_graph`): Converts a topological logical graph

    into a linear list of SASS instructions. Supports 1:1 opcode mapping via semantics and 1:N expansion via Kernel Macros (e.g. Conv2d loops).

    • Source Transformation (`to_python`): Converts SASS AST nodes back into Python LibCST nodes for high-level analysis or documentation.

  2. SassBackend: The CompilerBackend adapter for the Registry.

Attributes

MAX_REGISTERS

Classes

RegisterAllocator

Manages the mapping between symbolic variable names and physical registers.

SassSynthesizer

Bidirectional transpiler component.

SassBackend

Compiler Backend implementation for NVIDIA SASS.

Module Contents

ml_switcheroo.compiler.backends.sass.synthesizer.MAX_REGISTERS = 255
class ml_switcheroo.compiler.backends.sass.synthesizer.RegisterAllocator[source]

Manages the mapping between symbolic variable names and physical registers.

Implements a simple linear allocation strategy (bump pointer). A more complex implementation would handle liveness analysis and register spilling, but this suffices for basic block translation.

get_register(var_name: str) → ml_switcheroo.compiler.frontends.sass.nodes.Register[source]

Retrieves or allocates a register for a symbolic variable.

If the variable has been seen before, returns the existing register mapping. If new, allocates the next available physical register.

Parameters:

var_name (str) – The logical identifier (e.g., ‘input_1’, ‘bias’).

Returns:

A populated Register node (e.g., Register(‘R0’)).

Return type:

Register

Raises:

ValueError – If the allocator runs out of physical registers (>255).

allocate_temp() → ml_switcheroo.compiler.frontends.sass.nodes.Register[source]

Allocates a temporary anonymous register.

Useful for intermediate calculations or immediate loading.

Returns:

A new physical register.

Return type:

Register

reset() → None[source]

Resets the allocator state.

class ml_switcheroo.compiler.backends.sass.synthesizer.SassSynthesizer(semantics: ml_switcheroo.semantics.manager.SemanticsManager)[source]

Bidirectional transpiler component.

Handles: 1. Forward (Graph -> SASS): Synthesizes Assembly from Logical Graphs.

Delegates high-level ops (Conv2d, Linear) to Macros, and low-level ops (Add, Mul) to Semantic Opcode Lookup.

  1. Reverse (SASS -> Python): Synthesizes Python AST from Assembly nodes.

semantics
allocator
macro_registry: Dict[str, Callable]
from_graph(graph: ml_switcheroo.compiler.ir.LogicalGraph) → List[ml_switcheroo.compiler.frontends.sass.nodes.SassNode][source]

Converts a LogicalGraph into a list of SASS AST nodes.

Process: 1. Sorts nodes topologically. 2. Traverses nodes. 3. For each node:

  1. Check if it matches a Macro (e.g. Conv2d). If so, expand kernel.

  2. If not, lookup abstract opcode mapping (e.g. Add -> FADD).

  3. Allocate/Resolve Input Registers.

  4. Allocate Output Register.

  5. Construct Instruction node.

  1. Handles Input nodes by pre-allocating registers (Contract: R0, R1…).

Parameters:

graph (LogicalGraph) – The input computation graph.

Returns:

A structured list of assembly nodes.

Return type:

List[SassNode]

to_python(sass_nodes: List[ml_switcheroo.compiler.frontends.sass.nodes.SassNode]) → libcst.Module[source]

Converts SASS AST nodes into a Python source structure representation.

Used for analysis or round-trip verification. Registers are treated as variables. Instructions map to function calls sass.OPCODE(args).

Structure:

R0 = sass.FADD(R1, R2)

Parameters:

sass_nodes (List[SassNode]) – List of parsed SASS nodes.

Returns:

A LibCST module containing the Python representation.

Return type:

cst.Module

class ml_switcheroo.compiler.backends.sass.synthesizer.SassBackend(semantics: ml_switcheroo.semantics.manager.SemanticsManager | None = None)[source]

Bases: ml_switcheroo.compiler.backend.CompilerBackend

Compiler Backend implementation for NVIDIA SASS. Orchestrates the synthesis (Graph -> AST) and emission (AST -> Text).

synthesizer
emitter
compile(graph: ml_switcheroo.compiler.ir.LogicalGraph) → str[source]

Compiles LogicalGraph to SASS Assembly string.

Parameters:

graph – The intermediate representation.

Returns:

The SASS code.

Return type:

str