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 int_literal(self, args):
return Literal(int, int(args[0]))
def float_literal(self, args):
return Literal(float, float(args[0]))