Class IRBuilder
- Namespace
- AiDotNet.JitCompiler
- Assembly
- AiDotNet.dll
Builds an IR graph from a ComputationNode graph.
public class IRBuilder
- Inheritance
-
IRBuilder
- Inherited Members
Remarks
The IRBuilder converts a high-level ComputationNode graph (produced by autodiff) into a low-level IR graph suitable for optimization and compilation. It traverses the computation graph, converts each node to an IR operation, and builds the complete IR representation.
For Beginners: This translates autodiff graphs into a form the JIT compiler can work with.
Think of it like translating a recipe:
- Input: ComputationNode graph (high-level description of what to compute)
- Output: IR graph (low-level description ready for optimization)
The IRBuilder:
- Walks through all the computation nodes
- Identifies what operation each node represents
- Creates corresponding IR operations
- Builds a complete IR graph with inputs, operations, and outputs
This IR graph can then be optimized and compiled to fast executable code.
Methods
BuildBackward<T>(ComputationNode<T>, List<ComputationNode<T>>)
Builds a backward IR graph for gradient computation.
public IRGraph BuildBackward<T>(ComputationNode<T> outputNode, List<ComputationNode<T>> inputs)
Parameters
outputNodeComputationNode<T>The output node of the forward computation graph.
inputsList<ComputationNode<T>>The input nodes to compute gradients for.
Returns
- IRGraph
An IR graph that computes gradients via backpropagation.
Type Parameters
TThe numeric type used in the computation.
Remarks
This method builds the backward pass (gradient computation) graph from a forward graph. The backward graph takes output gradients as inputs and computes gradients with respect to the original inputs via automatic differentiation.
For Beginners: This creates the gradient computation graph for training.
In neural network training:
- Forward pass: input → layers → output → loss
- Backward pass: loss gradient → layers (in reverse) → input gradients
This method creates the backward pass graph automatically!
Algorithm:
- Traverse forward graph in reverse topological order
- For each operation, generate its backward (gradient) operation
- Handle gradient accumulation for nodes with multiple consumers
- Build IR graph mapping output gradients → input gradients
Example operations and their gradients:
- Add(a, b) → backward distributes gradient to both a and b
- MatMul(a, b) → backward: grad_a = grad_out @ b^T, grad_b = a^T @ grad_out
- ReLU(x) → backward: grad_x = grad_out * (x > 0)
IMPLEMENTATION STATUS:
This is a complex feature requiring implementation of:
Reverse Graph Traversal
- Walk forward graph in reverse topological order
- Track gradient flow through each operation
Backward Operation Mapping
- For each forward op type, generate corresponding backward op(s)
- Examples:
- AddOp → GradAddOp (distributes gradient to both inputs)
- MatMulOp → GradMatMulLeftOp + GradMatMulRightOp
- ReLUOp → GradReLUOp (masks gradient by activation)
- Etc. for all 43+ operation types
Gradient Accumulation
- When a node has multiple consumers, accumulate gradients
- Insert GradAccumulateOp to sum gradients from different paths
Memory Optimization
- Forward activations may need to be saved for backward pass
- Implement checkpointing for memory-efficient training
IR Operation Types Needed
- Create new IR op types for backward operations:
- GradAddOp, GradSubtractOp, GradMultiplyOp
- GradMatMulLeftOp, GradMatMulRightOp
- GradReLUOp, GradSigmoidOp, GradTanhOp
- GradConv2DOp, GradMaxPool2DOp
- GradAccumulateOp (sums multiple gradients)
- Implement code generation for each
- Create new IR op types for backward operations:
Testing Required
- Gradient correctness tests (numerical gradient checking)
- Performance benchmarks vs. non-compiled backward pass
- Memory usage profiling
TODO: Full implementation of backward pass IR builder
- This is a substantial feature requiring:
- New IR operation types (~50+ backward ops)
- Code generation for backward ops
- Gradient accumulation logic
- Extensive testing
- Estimated effort: 1-2 weeks for complete implementation
- See PyTorch's autograd and TensorFlow's GradientTape for reference implementations
Exceptions
- NotImplementedException
This method requires full implementation of backward operation mapping and gradient accumulation.
Build<T>(ComputationNode<T>, List<ComputationNode<T>>)
Builds an IR graph from a ComputationNode graph.
public IRGraph Build<T>(ComputationNode<T> outputNode, List<ComputationNode<T>> inputs)
Parameters
outputNodeComputationNode<T>The output node of the computation graph.
inputsList<ComputationNode<T>>The input nodes to the computation graph.
Returns
- IRGraph
An IR graph representing the computation.
Type Parameters
TThe numeric type used in the computation.
Remarks
This method performs a topological traversal of the computation graph, converting each ComputationNode to an IROp and building the complete IR graph. It handles input mapping, operation conversion, and output identification.
For Beginners: This converts a computation graph to IR format.
The process:
- Identifies all input nodes and assigns them tensor IDs
- Traverses the graph in topological order (inputs to outputs)
- Converts each node to an IR operation
- Builds the final IR graph with all operations connected
Example: If you have a graph: result = ReLU(MatMul(input, weights) + bias) This will create an IR graph with:
- Input tensors: input (t0), weights (t1), bias (t2)
- Operations: MatMul (t3 = MatMul(t0, t1)), Add (t4 = Add(t3, t2)), ReLU (t5 = ReLU(t4))
- Output: t5
Exceptions
- InvalidOperationException
Thrown if a node doesn't have operation type metadata or uses an unsupported operation.