Support symbolic arithmetic operations #2736
Description
🚀 Feature
Thunder JIT needs to support arithmetic operations on symbolic values, producing new symbolic expressions. This includes addition, subtraction, and composition of symbolic values.
Current Behavior
Arithmetic on dynamic values likely forces eager evaluation and concretizing the tracing time results.
Expected Behavior
Arithmetic operations on symbolic values should produce symbolic results:
Sym(s0) + 1→Sym(s0 + 1)Sym(s0) - 1056→Sym(s0 - 1056)Sym(s0) + Sym(s1)→Sym(s0 + s1)
Example from torch.compile
add: "Sym(s67 + 1)" = l_kwargs_past_key_values_layers_0_cumulative_length + 1
add: "Sym(s50 + s67)" = l_kwargs_past_key_values_layers_0_cumulative_length + s50
add_10: "Sym(s50 + s59)" = l_kwargs_past_key_values_layers_1_cumulative_length + s50Minimal Reproduction Case
import torch
import thunder
@thunder.jit
def symbolic_arithmetic(cumulative_length: int, seq_len: int) -> tuple[int, int, int]:
"""Both inputs should be symbolic."""
# Pattern 1: Symbolic + constant
new_length_1 = cumulative_length + 1
# Pattern 2: Symbolic - constant
offset = cumulative_length - 1056
# Pattern 3: Symbolic + symbolic
combined = cumulative_length + seq_len
return new_length_1, offset, combined
result = symbolic_arithmetic(1024, 32)
assert result == (1025, -32, 1056)
print(symbolic_arithmetic._lc_cs.last_epilogue_traces[-1])
# This shows that the values are concretized and not treated symbolically
# def epilogue():
# return (1025, -32, 1056)Operations to Support
- Addition:
symbolic + constant,symbolic + symbolic - Subtraction:
symbolic - constant,symbolic - symbolic - Multiplication:
symbolic * constant - Floor division:
symbolic // constant
Success Criteria
- All arithmetic operations preserve symbolic types
- Symbolic expressions can be composed
- Generated code correctly handles runtime values
- Works with both scalar-scalar and scalar-tensor operations