def single_loop_shape(self, args):
expr = args[0]
loop_var = expr.pprint()
default_greater_eq = Expr([Action(None, Access(Var(f'{loop_var}_greater_eq')))])
default_less_eq = Expr([Action(None, Access(Var(f'{loop_var}_less_eq')))])
default_step = Expr([Action(None, Literal(int, 1))])
return LoopShape(expr.clone(),
default_greater_eq,
default_less_eq,
default_step)
def loop_shape_parts(self, args):
merged = None
for loop_shape_builder in args:
if merged is None:
merged = loop_shape_builder
else:
merged.merge(loop_shape_builder)
assert (merged is not None)
assert (merged.loop_var is not None)
loop_var = merged.loop_var.pprint()
default_greater_eq = Expr(
[Action(None, Access(Var(f'{loop_var}_greater_eq')))])
default_less_eq = Expr(
[Action(None, Access(Var(f'{loop_var}_less_eq')))])
default_step = Expr([Action(None, Literal(int, 1))])
return merged.build(default_greater_eq, default_less_eq, default_step)
def array_access(self, args):
return Access(args[0], args[1:])
def scalar_access(self, args):
return Access(args[0])