-
Notifications
You must be signed in to change notification settings - Fork 70
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
transformations: (csl) added transformation for builtin constructs to csl #2713
Conversation
This transformation converts func.func and func.return to their csl counterparts, wraps everything in a program module and adds a layout module.
This depends on #2712 to make |
@dataclass | ||
class AddLayout(RewritePattern): | ||
pe_program: str | ||
pe_x_count: int | ||
pe_y_count: int | ||
launch_color: int | ||
ctx: MLContext | ||
|
||
def _build_layout_module(self): | ||
return Parser( | ||
self.ctx, | ||
f""" | ||
"csl.module"() <{{kind = #csl<module_kind layout>}}> ({{ | ||
|
||
%LAUNCH = "csl.get_color"() <{{id = {self.launch_color} : i5}}> : () -> !csl.color | ||
|
||
%memcpy_init_params = "csl.const_struct"(%LAUNCH) <{{ | ||
items = {{ width = {self.pe_x_count} : i32, height = {self.pe_y_count} : i32}}, | ||
ssa_fields = ["LAUNCH"] | ||
}}> : (!csl.color) -> !csl.comptime_struct | ||
|
||
%memcpy = "csl.import_module"(%memcpy_init_params) <{{module = "<memcpy/get_params>"}}> : (!csl.comptime_struct) -> !csl.imported_module | ||
|
||
csl.layout {{ | ||
|
||
%x_dim_idx = arith.constant {self.pe_x_count} : index | ||
%y_dim_idx = arith.constant {self.pe_x_count} : index | ||
%x_dim = arith.index_cast %x_dim_idx : index to i32 | ||
%y_dim = arith.index_cast %y_dim_idx : index to i32 | ||
|
||
"csl.set_rectangle"(%x_dim, %y_dim) : (i32, i32) -> () | ||
|
||
%c0 = arith.constant 0 : index | ||
%c1 = arith.constant 1 : index | ||
scf.for %x_coord_idx = %c0 to %x_dim_idx step %c1 {{ | ||
scf.for %y_coord_idx = %c0 to %y_dim_idx step %c1 {{ | ||
%x_coord = arith.index_cast %x_coord_idx : index to i32 | ||
%y_coord = arith.index_cast %y_coord_idx : index to i32 | ||
%memcpy_params = "csl.member_call"(%memcpy, %x_coord) <{{field = "get_params"}}> : (!csl.imported_module, i32) -> !csl.comptime_struct | ||
%tile_code_params = "csl.const_struct"(%memcpy_params) <{{ssa_fields = ["memcpy_params"]}}> : (!csl.comptime_struct) -> !csl.comptime_struct | ||
"csl.set_tile_code"(%x_coord, %y_coord, %tile_code_params) <{{file = "{self.pe_program}"}}> : (i32, i32, !csl.comptime_struct) -> () | ||
}} | ||
}} | ||
|
||
}} | ||
}}) {{sym_name = "layout"}} : () -> () | ||
""", | ||
).parse_op() | ||
|
||
@op_type_rewrite_pattern | ||
def match_and_rewrite(self, op: csl.CslModuleOp, rewriter: PatternRewriter, /): | ||
rewriter.insert_op(self._build_layout_module(), InsertPoint.after(op)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Not sure this is the right approach, let's chat once we have the baseline testfiles ready! 😄
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am not sure either, TBH. Happy to discuss alternatives! This will need fixing anyway after all the other stuff is merged.
I don't think this is needed anymore, so I'll close it. It can be reopened if that's not the case |
This transformation converts
func.func
andfunc.return
to theircsl
counterparts, wraps everything in a program module and adds a layout module.