def single_loop_shape(self, args): loop_var = args[0].var default_greater_eq = Access(f'{loop_var}_greater_eq') default_less_eq = Access(f'{loop_var}_less_eq') default_step = Literal(int, 1) return LoopShape(args[0], 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.var default_greater_eq = Access(f'{loop_var}_greater_eq') default_less_eq = Access(f'{loop_var}_less_eq') default_step = 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])