burn-ir

pub enum OperationIr {
    /// Basic operation on a float tensor.
    BaseFloat(BaseOperationIr),
    /// Basic operation on an int tensor.
    BaseInt(BaseOperationIr),
    /// Basic operation on a bool tensor.
    BaseBool(BaseOperationIr),
    /// Numeric operation on a float tensor.
    NumericFloat(DType, NumericOperationIr),
    /// Numeric operation on an int tensor.
    NumericInt(DType, NumericOperationIr),
    /// Operation specific to a bool tensor.
    Bool(BoolOperationIr),
    /// Operation specific to an int tensor.
    Int(IntOperationIr),
    /// Operation specific to a float tensor.
    Float(DType, FloatOperationIr),
    /// Module operation.
    Module(ModuleOperationIr),
    /// Initialize operation.
    Init(InitOperationIr),
    /// A custom operation.
    Custom(CustomOpIr),
    /// A tensor is dropped.
    Drop(TensorIr),
    /// Operation specific to a distributed tensor.
    Distributed(DistributedOperationIr),
    /// Activation function operation (relu, gelu, sigmoid, softmax, …).
    Activation(ActivationOperationIr),
}

pub struct BinaryOpIr {
    pub lhs: TensorIr,
    pub rhs: TensorIr,
    pub out: TensorIr,
}

impl_ir_create!(
    MatmulOpIr {
        lhs: TensorIr,
        rhs: TensorIr
    },
    shape = calculate_matmul_output(&lhs.shape, &rhs.shape).unwrap(),
    dtype = output_dtype_mixed([&lhs.dtype, &rhs.dtype]).unwrap(),
    // Additional constructor for mixed dtypes
    create_mixed(out_dtype: DType)
);

struct Once(AtomicBool);
impl FusionObserver for Once {
  fn registered(&self, op: &OperationIr) {
    if matches!(op, OperationIr::NumericFloat(..)) && !self.0.swap(true, Ordering::SeqCst) {
      println!("{op:#?}");
      println!("json: {}", serde_json::to_string(op).unwrap());
    }
  }
  fn block_starts(&self, _block: BlockOperations<'_>) {}
  fn block_ran(&self) {}
}

let _obs = FusionObservation::new(Arc::new(Once(AtomicBool::new(false))));
let a = small(&device, 1.0); // 2x2 Tensor::from_floats on Device::metal(..)
let b = small(&device, 2.0);
let d = a * b;
let _ = d.into_data();

NumericFloat(
    F32,
    Mul(
        BinaryOpIr {
            lhs: TensorIr {
                id: TensorId {
                    value: 1,
                },
                shape: Shape {
                    dims: [
                        2,
                        2,
                    ],
                },
                status: ReadWrite,
                dtype: F32,
            },
            rhs: TensorIr {
                id: TensorId {
                    value: 2,
                },
                ...
                status: ReadWrite,
                dtype: F32,
            },
            out: TensorIr {
                id: TensorId {
                    value: 3,
                },
                ...
                status: NotInit,
                dtype: F32,
            },
        },
    ),
)
json: {"NumericFloat":["F32",{"Mul":{"lhs":{"id":{"value":1},"shape":{"dims":[2,2]},"status":"ReadWrite","dtype":"F32"},"rhs":{"id":{"value":2},"shape":{"dims":[2,2]},"status":"ReadWrite","dtype":"F32"},"out":{"id":{"value":3},"shape":{"dims":[2,2]},"status":"NotInit","dtype":"F32"}}}]}

pub struct TensorIr {
    /// The [tensor id](TensorId).
    pub id: TensorId,
    /// The shape of the tensor.
    pub shape: Shape,
    /// The [status](TensorStatus) of the tensor when it was used.
    pub status: TensorStatus,
    /// The [type](DType) of the tensor.
    pub dtype: DType,
}

pub struct HandleContainer<H> {
    handles: HashMap<TensorId, Handle<H>>,
    counter: u64,
    // ...
    errored: usize,
}

        let (id, handle) = self
            .take(id)
            .unwrap_or_else(|| panic!("Should have handle for tensor {id:?}"));

        match handle {
            Handle::Existing(handle) => match status {
                TensorStatus::ReadOnly => {
                    self.put(id, Handle::Existing(handle.clone()));
                    handle
                }
                TensorStatus::ReadWrite => handle,
                TensorStatus::NotInit => panic!(
                    "Cannot get uninitialized tensor {id:?}. Tensor exist but with wrong status"
                ),
            },

fn compute(device: &Device) -> Tensor<2> {
  let a = Tensor::<2>::from_floats([[1.0, -2.0], [3.0, -4.0]], device);
  let b = Tensor::<2>::from_floats([[5.0, 6.0], [7.0, 8.0]], device);
  relu(a * b + 1.0)
}

// Metal: a FusionObserver prints every registered op.
let out = compute(&metal);

// Capture: the scope records instead of executing.
let cap = Device::capture();
let graph = cap
  .capture_scope(|scope| {
    let _out = compute(&cap);
    scope.complete([], [])
  })
  .unwrap();
for op in &graph.graph.operations { /* print the variant */ }

--- Metal (Fusion)
  fusion registered: Init
  fusion registered: Init
  fusion registered: NumericFloat(F32, Mul)
  fusion registered: NumericFloat(F32, AddScalar)
  fusion registered: NumericFloat(F32, LowerEqualElem)
  fusion registered: BaseFloat(MaskFill)
  result [6.0, 0.0, 22.0, 0.0]
--- Device::capture()
  captured: Init
  captured: Init
  captured: NumericFloat(F32, Mul)
  captured: NumericFloat(F32, AddScalar)
  captured: Activation(Relu)
  inputs [] outputs [], 2 constant(s)
  GraphIr as JSON: 924 bytes

  [registered] Init in=[] out=[t1:NotInit]
  ...
  [registered] Init in=[] out=[t2:NotInit]
  ...
  [registered] Init in=[] out=[t3:NotInit]
  ...
  [registered] NumericFloat(F32, Mul) in=[t1:ReadWrite, t2:ReadWrite] out=[t4:NotInit]
  [registered] NumericFloat(F32, Add) in=[t4:ReadWrite, t3:ReadWrite] out=[t5:NotInit]
  [registered] Float(F32, Exp) in=[t5:ReadWrite] out=[t6:NotInit]