mitsuba-renderer / drjit

Dr.Jit — A Just-In-Time-Compiler for Differentiable Rendering
BSD 3-Clause "New" or "Revised" License
603 stars 45 forks source link

Handle list comprehensions in `@dr.syntax` #256

Closed njroussel closed 3 months ago

njroussel commented 3 months ago

This PR fixes #252.

List comprehensions and other generators define statement-local variables that can never be part of the state loop. This PR guarantees this by ignoring those variables names during the AST traversal of @dr.syntax.

I've added a test that isn't much more than a regressions test - it doesn't actually verify that the generator's variables become part of the loop state.

I think this is a fairly safe change @wjakob, but I'm also a bit adverse to introducing more and more handling of these edge cases in @dr.syntax

wjakob commented 3 months ago

Hi @njroussel,

thank you for this fix. I had not considered comprehensions when designing the AST transformations. It's a natural Python construct that should most definitely be supported.

But I'm thinking that the implementation could potentially be even easier.

First, comp.target.ctx can AFAIK only be an ast.Store, so I'm thinking that the two cases aren't needed.

Second, you cannot make assignments to outside variables in a comprehension. So at the end, you can just set self.var_w = var_w. In fact, maybe nothing at all needs to be done with self.var_w (no copies, no re-assignment) since we know that it won't be changed.

Last, it is important when a comprehension reads additional variables, but we can exclude variables that the comprehension sets itself in its isolated scope.

How about self.var_r = (self.var_r - comp_targets_w) | var_r (actually, I think that's what you had, just on one line).

njroussel commented 3 months ago

Indeed that simplified the code a bit, thanks.