ml_switcheroo.compiler.backends.rdna.macros¶

RDNA Macro Expansion Logic.

Procedural generators for complex RDNA kernel logic (Conv2d, Linear).

Classes¶

RegisterAllocatorProtocol

Protocol for the Dual-Pool Register Allocator.

Functions¶

expand_conv2d(...)

Generates the RDNA assembly kernel for a 2D Convolution loop.

expand_linear(...)

Generates the RDNA assembly kernel for a Linear Layer.

Module Contents¶

class ml_switcheroo.compiler.backends.rdna.macros.RegisterAllocatorProtocol[source]¶

Bases: Protocol

Protocol for the Dual-Pool Register Allocator.

get_vector_register(var_name: str) → ml_switcheroo.compiler.frontends.rdna.nodes.VGPR[source]¶

Gets or allocates a VGPR.

get_scalar_register(var_name: str) → ml_switcheroo.compiler.frontends.rdna.nodes.SGPR[source]¶

Gets or allocates an SGPR.

allocate_vector_temp() → ml_switcheroo.compiler.frontends.rdna.nodes.VGPR[source]¶

Allocates temp VGPR.

allocate_scalar_temp() → ml_switcheroo.compiler.frontends.rdna.nodes.SGPR[source]¶

Allocates temp SGPR.

ml_switcheroo.compiler.backends.rdna.macros.expand_conv2d(allocator: RegisterAllocatorProtocol, node_id: str, metadata: Dict[str, Any]) → List[ml_switcheroo.compiler.frontends.rdna.nodes.RdnaNode][source]¶

Generates the RDNA assembly kernel for a 2D Convolution loop.

ml_switcheroo.compiler.backends.rdna.macros.expand_linear(allocator: RegisterAllocatorProtocol, node_id: str, metadata: Dict[str, Any]) → List[ml_switcheroo.compiler.frontends.rdna.nodes.RdnaNode][source]¶

Generates the RDNA assembly kernel for a Linear Layer.