Fix sizes of literals in frozen functions

[dvicini]

This addresses #498

This fix was AI generated, so it needs careful review. The core ideas seem reasonable, but I am not aware of possible pitfalls.

[DoeringChristian]

Hi @dvicini,

Sorry for the delayed response.
In principle this looks fine. However, we are essentially trading off two different failure conditions. Multiplying the value by 0.0 will simplify it to a literal internally. We therefore loose the information about which input variable the literal size is related to. Adding the literal sizes to the variable the equivalence classes would solve this particular problem, as now the output literal gets tracked as having the same size as the input variable. However, this could introduce other bugs:

import drjit as dr

Float = dr.llvm.Float

def func(x):
    fixed = dr.ones(Float, 3)
    return fixed

frozen = dr.freeze(func)

# Recording: x=3, literal=3 -> same size class
x = Float([1, 2, 3])
res = frozen(x)
ref = func(x)
print(f"record: res_w={dr.width(res)} ref_w={dr.width(ref)}")

# Replay: x=5 -> literal should stay 3, not become 5
x2 = Float([1, 2, 3, 4, 5])
res2 = frozen(x2)
ref2 = func(x2)
print(f"replay: res_w={dr.width(res2)} ref_w={dr.width(ref2)} PASS={dr.width(res2)==dr.width(ref2)}")

Because we fundamentally loose the information about input-output relations, for literals, we could decide for one or the other approach. I'm fine with it either way (not sure what causes more confusion)? I tested your implementation for performance. The fix introduces a bunch of hash-set lookups for literals, which make up the largest part of the variables in most scenes. It it seems fine (slight regression on bistro scene, but within noise for my tests):

  ┌───────────────┬───────────────────┬───────────────────┐                                                                                                                                                          
  │               │   Fix disabled    │    Fix enabled    │                                                                                                                                                          
  ├───────────────┼───────────────────┼───────────────────┤                                                                                                                                                          
  │ Replay avg    │ 0.0467s ± 0.0029s │ 0.0479s ± 0.0030s │       
  ├───────────────┼───────────────────┼───────────────────┤
  │ Replay min    │ 0.0401s           │ 0.0411s           │                                                                                                                                                          
  ├───────────────┼───────────────────┼───────────────────┤
  │ Speedup (avg) │ 60.04x            │ 58.61x            │                                                                                                                                                          
  ├───────────────┼───────────────────┼───────────────────┤                                                                                                                                                          
  │ No-freeze avg │ 2.8065s ± 0.0232s │ 2.8051s ± 0.0316s │
  ├───────────────┼───────────────────┼───────────────────┤                                                                                                                                                          
  │ Record        │ 2.8504s           │ 2.8713s           │       
  └───────────────┴───────────────────┴───────────────────┘                                                                                                                                                          

(measured by profiling 10 times).

Best
Christian

[dvicini]

Hi Christian,

I see, the size suddenly becoming "5" in your example also seems very confusing. My fix was purely AI generated, so there is no need to use it as is, if it's not the best solution. Isn't it somehow possible in this test to ensure that the literal stays of the size 3, at which it was explicitly created?

I am not too concerned about the performance hit, as long as we can get correct behavior for various versions of this test.