split convert_stream_to_snax_stream in scheduler and layout resolution #332
Conversation
There was a problem hiding this comment.
I have some small organizational proposed changes.
accfg registry is definitly a bit messy in this regard, which is a shame, but also no deal-breaker right now.
Can you also maybe update the comments a bit?
I guess more changes are also coming in future PRs, so please do when you see fit!
|
|
||
|
|
||
| @dataclass | ||
| class MemrefStreamToSnaxPattern(RewritePattern): | ||
| class AutoflowScheduler(RewritePattern): |
There was a problem hiding this comment.
Can you update the documentation here?
| # get template and template_bounds | ||
| accelerator_type = AcceleratorRegistry().get_acc_info(acc_op) | ||
| assert issubclass(accelerator_type, SNAXStreamer) | ||
|
|
||
| template = accelerator_type.get_template(op) |
There was a problem hiding this comment.
Excited to move this into the IR at some point!
| from compiler.ir.stream.access_pattern import Template | ||
|
|
||
|
|
||
| def get_accelerator_info(op: stream.StreamingRegionOpBase) -> Template: |
There was a problem hiding this comment.
| def get_accelerator_info(op: stream.StreamingRegionOpBase) -> Template: | |
| def get_accelerator_template(op: stream.StreamingRegionOpBase) -> Template: |
There was a problem hiding this comment.
Can you make this a method of the (base class of the) op itself maybe?
There was a problem hiding this comment.
unfortunately this would be a circular import
| class LayoutResolution(RewritePattern): | ||
| @op_type_rewrite_pattern | ||
| def match_and_rewrite(self, op: stream.ScheduleOp, rewriter: PatternRewriter): | ||
| template = get_accelerator_info(op) |
There was a problem hiding this comment.
If this were a method:
| template = get_accelerator_info(op) | |
| template = op.get_accelerator_template() |
jorendumoulin
force-pushed
the
joren/split-pass
branch
from
b1d9213 to
09c31da
Compare
stacked on #331
With the addition of the schedule op, this PR splits the
convert_stream_to_snax_stream(which converts operation -> access pattern) into two separate passes (scheduler: operation -> schedule), (layout resolution: schedule -> access pattern).The schedule representation is nice to have to run other passes on before doing layout resolution, such as selecting a layout.