gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] For loops now promote initial value (#524)

Browse Source
This commit is contained in:
Philippe Tillet
2022-05-24 13:20:10 -07:00
committed by GitHub
parent 96bff90471
commit 011bc83c1b
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/triton/code_gen.py
View File

@@ -584,13 +584,17 @@ class CodeGenerator(ast.NodeVisitor):
for stmt in node.orelse: for stmt in node.orelse:
ast.NodeVisitor.generic_visit(self, stmt) ast.NodeVisitor.generic_visit(self, stmt)
return return
# create nodes # create nodes
st_target = ast.Name(id=node.target.id, ctx=ast.Store()) st_target = ast.Name(id=node.target.id, ctx=ast.Store())
ld_target = ast.Name(id=node.target.id, ctx=ast.Load()) ld_target = ast.Name(id=node.target.id, ctx=ast.Load())
arg_0 = node.iter.args[0] if len(node.iter.args) > 1 else ast.Num(0) arg_0 = node.iter.args[0] if len(node.iter.args) > 1 else ast.Num(0)
arg_1 = node.iter.args[1] if len(node.iter.args) > 1 else node.iter.args[0] arg_1 = node.iter.args[1] if len(node.iter.args) > 1 else node.iter.args[0]
arg_2 = node.iter.args[2] if len(node.iter.args) > 2 else ast.Num(1) arg_2 = node.iter.args[2] if len(node.iter.args) > 2 else ast.Num(1)
# init node
init_node = ast.Assign(targets=[st_target], value=arg_0) init_node = ast.Assign(targets=[st_target], value=arg_0)
# step node
pos_cond_node = ast.Compare(ld_target, [ast.Lt()], [arg_1]) pos_cond_node = ast.Compare(ld_target, [ast.Lt()], [arg_1])
neg_cond_node = ast.Compare(ld_target, [ast.Gt()], [arg_1]) neg_cond_node = ast.Compare(ld_target, [ast.Gt()], [arg_1])
pos_step_node = ast.Compare(arg_2, [ast.Gt()], [ast.Num(0)]) pos_step_node = ast.Compare(arg_2, [ast.Gt()], [ast.Num(0)])
@@ -610,7 +614,17 @@ class CodeGenerator(ast.NodeVisitor):
cond = build_cond() cond = build_cond()
return self.builder.cond_br(cond.handle, loop_bb, next_bb) return self.builder.cond_br(cond.handle, loop_bb, next_bb)
# init loop induction variable
self.visit(init_node) self.visit(init_node)
# promote it to right type
init_val = self.value_constructor.get_value(node.target.id)
promote = lambda a, b: triton.language.semantic.computation_type_impl(a, b, False)
start_ty = triton.language.core._to_tensor(iter_args[0], self.builder).type
stop_ty = triton.language.core._to_tensor(iter_args[1], self.builder).type if len(iter_args) > 1 else None
ty = promote(start_ty, stop_ty) if len(iter_args) > 1 else start_ty
casted = triton.language.semantic.cast(init_val, ty, self.builder)
self.value_constructor.set_value(node.target.id, casted)
# create cond
cond = build_cond() cond = build_cond()
self.builder.cond_br(cond.handle, loop_bb, next_bb) self.builder.cond_br(cond.handle, loop_bb, next_bb)
self.builder.set_insert_block(loop_bb) self.builder.set_insert_block(loop_bb)