feat[next]: Prepare TraceShift pass for GTIR #1592
Conversation
tehrengruber
changed the title
|
|
||
| return recorded_shifts | ||
| trace_stencil = TraceShifts.trace_stencil |
There was a problem hiding this comment.
Why is it also a classmethod?
There was a problem hiding this comment.
To allow subclassing.
There was a problem hiding this comment.
Did you just make this up or do you have a use-case in mind?
There was a problem hiding this comment.
To stick to your wording: I made this up. I can change it if you like, I can also make up unlikely cases where this might be useful.
| instance = cls() | ||
| instance.visit(node) | ||
| ctx = instance.initialize_context(args) | ||
| instance.visit(im.call(stencil)(*args), ctx=ctx) |
There was a problem hiding this comment.
This pattern looks a bit weird as initialization is done in 2 steps. Maybe you can disentangle a bit more?
There was a problem hiding this comment.
I cleaned it up a bit, but still not great. I would leave it as is for now.
| if isinstance(stencil, ir.Lambda): | ||
| assert num_args is None or num_args == len(stencil.params) | ||
| num_args = len(stencil.params) |
There was a problem hiding this comment.
what's special with lambdas? what else can it be?
There was a problem hiding this comment.
For lambdas we can deduce the number of arguments and the num_args parameter is optional. This is essentially a convenience feature for testing. I've added a comment.
|
|
||
| return recorded_shifts | ||
| trace_stencil = TraceShifts.trace_stencil |
There was a problem hiding this comment.
Did you just make this up or do you have a use-case in mind?
The pass now only operates on stencils instead of
StencilClosures using the newly introduced functiontrace_stencil.