Enzyme Dialect
Reactant.MLIR.Dialects.enzyme.binomial_progress Method
binomial_progress
Given num_steps remaining iterations and a budget of available checkpoints, returns how many steps to advance before placing the next checkpoint. This is the classic Revolve "split" function used for binomial checkpointing of loops.
Writing beta(s, t) = C(s + t, t) for the longest chain reversible with s checkpoints and at most t recomputations of each step, and taking t minimal with beta(budget, t) >= num_steps, every advance in
[ num_steps - beta(budget-1, t) , beta(budget, t-1) ]attains that optimal t, and the interval is non-empty by Pascal's rule beta(s,t) = beta(s-1,t) + beta(s,t-1). This op returns the lower edge, clamped to [1, num_steps-1], which also lands within a few percent of the minimal total recomputation.
Two boundary cases matter to callers. num_steps <= 1 returns num_steps (so 0 or 1). budget <= 1 returns num_steps: with one checkpoint left the entire remaining stretch is replayed from it, so the advance covers all of it. That is also what makes the advances chosen for successive slots sum to exactly the trip count across budget slots, which lets a driver that iterates once per checkpoint still reach the end of the primal.
The result is an advance distance, not a repetition count: it grows like num_steps, not like num_steps^(1/budget).
num_steps and budget may be an index, a signless integer of any width, or a tensor thereof (with the result taking the same type). The tensor forms carry the scalar computation for dialects whose scalars are tensors: stablehlo callers use tensor<i64>, a rank-0 ranked tensor.
Reactant.MLIR.Dialects.enzyme.broadcast Method
broadcast
Broadcast the operand by adding extra dimensions with sizes provided by the shape attribute to the front. For scalar operands, ranked tensor is created.
NOTE: Only works for scalar and ranked tensor operands for now.
sourceReactant.MLIR.Dialects.enzyme.concat Method
concat
Concat list of input arguments into a generic value
sourceReactant.MLIR.Dialects.enzyme.dump Method
dump
Debug operation that dumps a tensor value with a label.
sourceReactant.MLIR.Dialects.enzyme.extract Method
extract
Extract value from batched operand at index
source