Table of Contents

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

outputNode ComputationNode<T>

The output node of the forward computation graph.

inputs List<ComputationNode<T>>

The input nodes to compute gradients for.

Returns

IRGraph

An IR graph that computes gradients via backpropagation.

Type Parameters

T

The 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:

  1. Traverse forward graph in reverse topological order
  2. For each operation, generate its backward (gradient) operation
  3. Handle gradient accumulation for nodes with multiple consumers
  4. 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:

  1. Reverse Graph Traversal

    • Walk forward graph in reverse topological order
    • Track gradient flow through each operation
  2. 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
  3. Gradient Accumulation

    • When a node has multiple consumers, accumulate gradients
    • Insert GradAccumulateOp to sum gradients from different paths
  4. Memory Optimization

    • Forward activations may need to be saved for backward pass
    • Implement checkpointing for memory-efficient training
  5. 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
  6. 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

outputNode ComputationNode<T>

The output node of the computation graph.

inputs List<ComputationNode<T>>

The input nodes to the computation graph.

Returns

IRGraph

An IR graph representing the computation.

Type Parameters

T

The 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:

  1. Identifies all input nodes and assigns them tensor IDs
  2. Traverses the graph in topological order (inputs to outputs)
  3. Converts each node to an IR operation
  4. 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.