Augmented assignment double-evaluates RHS for subscript and attribute targets #1233

Description

GH-1230 / commit 7af3419 fixed double evaluation of the RHS in augmented assignments for simple name targets (x += expr). However, the same bug still exists for subscript and attribute targets that fall back to make_new_assign_statement().

The root cause is in emit_AugAssign (codegen.py). rhs = adj.eval(node.value) evaluates the RHS expression at line 3416, but make_new_assign_statement() constructs a synthetic ast.Assign node containing the raw AST node.value, which gets re-evaluated when adj.eval(new_node) processes the BinOp:

def make_new_assign_statement():
    new_node = ast.Assign(targets=[lhs], value=ast.BinOp(lhs, node.op, node.value))
    adj.eval(new_node)  # re-evaluates node.value from raw AST

rhs = adj.eval(node.value)  # first evaluation (result unused by fallback paths)

Affected code paths

Every make_new_assign_statement() call site in emit_AugAssign double-evaluates:

1. Attribute targets (s.field += expr) — the ast.Attribute TODO path (line 3531)
2. Non-atomic array subscripts (arr_int16[i] += expr) — dtype in non_atomic_types (line 3433)
3. Composite non-atomic subscripts (arr_vec_uint8[i] += expr) — vec/mat/quat of non-atomic scalars (line 3445)
4. Unsupported operators on array subscripts (arr[i] *= expr, arr[i] **= expr) — no atomic implementation (lines 3484, 3506, 3523)
5. wp.adjoint[var] += expr in custom grad functions (line 3422)
6. Generic fallback for any other LHS type (line 3535)

Reproduction

All three scripts below print counter = 2 (expected 1):

import warp as wp

@wp.struct
class State:
    value: float

@wp.func
def side_effect_add(counter: wp.array(dtype=int), a: float, b: float) -> float:
    wp.atomic_add(counter, 0, 1)
    return a + b

@wp.kernel
def kernel(counter: wp.array(dtype=int), result: wp.array(dtype=float)):
    s = State()
    s.value = 1.0
    s.value += side_effect_add(counter, 2.0, 3.0)
    result[0] = s.value

counter = wp.zeros(1, dtype=int)
result = wp.zeros(1, dtype=float)
wp.launch(kernel, dim=1, inputs=[counter, result])
print(f"counter = {counter.numpy()[0]}  (expected 1)")
print(f"result  = {result.numpy()[0]}  (expected 6.0)")

import warp as wp

@wp.func
def side_effect_inc(counter: wp.array(dtype=int), val: wp.int16) -> wp.int16:
    wp.atomic_add(counter, 0, 1)
    return val

@wp.kernel
def kernel(counter: wp.array(dtype=int), data: wp.array(dtype=wp.int16)):
    i = wp.tid()
    data[i] += side_effect_inc(counter, wp.int16(1))

counter = wp.zeros(1, dtype=int)
data = wp.array([wp.int16(2)], dtype=wp.int16)
wp.launch(kernel, dim=1, inputs=[counter, data])
print(f"counter = {counter.numpy()[0]}  (expected 1)")
print(f"data[0] = {data.numpy()[0]}  (expected 3)")

import warp as wp

@wp.func
def side_effect_mul(counter: wp.array(dtype=int), val: float) -> float:
    wp.atomic_add(counter, 0, 1)
    return val

@wp.kernel
def kernel(counter: wp.array(dtype=int), data: wp.array(dtype=float)):
    i = wp.tid()
    data[i] *= side_effect_mul(counter, 2.0)

counter = wp.zeros(1, dtype=int)
data = wp.array([1.0], dtype=float)
wp.launch(kernel, dim=1, inputs=[counter, data])
print(f"counter = {counter.numpy()[0]}  (expected 1)")
print(f"data[0] = {data.numpy()[0]}  (expected 2.0)")

Proposed fix

The root cause is that make_new_assign_statement() conflates RHS evaluation with LHS assignment by constructing a synthetic AST that re-walks everything. The clean fix is to separate these concerns:

1. Extract an assign_to_target(lhs_node, value_var) method from emit_Assign that performs only the LHS dispatch and assignment, given an already-evaluated Var.

   emit_Assign currently does two things: (a) evaluate the RHS via adj.eval(node.value), then (b) dispatch on the LHS target type (ast.Name, ast.Subscript, ast.Attribute) to perform the store. Step (b) is exactly what emit_AugAssign's fallback paths need.

2. Replace make_new_assign_statement() entirely. Instead of constructing a synthetic AST, compute the binary operation result directly and call the new helper:

   # In emit_AugAssign, where make_new_assign_statement() is currently called:
   op_name = builtin_operators[type(node.op)]
   target = adj.eval(lhs)
   result = adj.add_builtin_call(op_name, [target, rhs])
   adj.assign_to_target(lhs, result)

3. Refactor emit_Assign to use the same helper:

   def emit_Assign(adj, node):
       # ... validation, RHS evaluation ...
       rhs = adj.eval(node.value)
       adj.assign_to_target(node.targets[0], rhs)

This eliminates double evaluation for all target types, removes the need to construct synthetic AST nodes, and makes the code easier to reason about — emit_AugAssign becomes "evaluate both sides once, compute the operation, assign the result" for every target type uniformly.