def visit_HorizontalIf(self, path: tuple, node_name: str,
node: gt_ir.HorizontalIf) -> Tuple[bool, gt_ir.If]:
assert self.extent is not None
conditions = []
for axis, interval in node.intervals.items():
extent = self.extent[self.domain.index(axis)]
if (interval.start.level == interval.end.level
and interval.start.offset == interval.end.offset - 1):
# Use a single condition
conditions.append(
gt_ir.BinOpExpr(
op=gt_ir.BinaryOperator.EQ,
lhs=gt_ir.AxisIndex(axis=axis),
rhs=gt_ir.AxisOffset(axis=axis,
endpt=interval.start.level,
offset=interval.start.offset),
))
else:
# start
if (interval.start.level != gt_ir.LevelMarker.START
or interval.start.offset > extent[0]):
conditions.append(
gt_ir.BinOpExpr(
op=gt_ir.BinaryOperator.GE,
lhs=gt_ir.AxisIndex(axis=axis),
rhs=gt_ir.AxisOffset(axis=axis,
endpt=interval.start.level,
offset=interval.start.offset),
))
# end
if interval.end.level != gt_ir.LevelMarker.END or interval.end.offset < extent[
1]:
conditions.append(
gt_ir.BinOpExpr(
op=gt_ir.BinaryOperator.LT,
lhs=gt_ir.AxisIndex(axis=axis),
rhs=gt_ir.AxisOffset(axis=axis,
endpt=interval.end.level,
offset=interval.end.offset),
))
if conditions:
return (
True,
gt_ir.If(
condition=functools.reduce(
lambda x, y: gt_ir.BinOpExpr(
op=gt_ir.BinaryOperator.AND, lhs=x, rhs=y),
conditions,
),
main_body=node.body,
),
)
else:
return True, node.body
def visit_BinOp(self, node: ast.BinOp) -> gt_ir.BinOpExpr:
op = self.visit(node.op)
rhs = gt_ir.utils.make_expr(self.visit(node.right))
lhs = gt_ir.utils.make_expr(self.visit(node.left))
result = gt_ir.BinOpExpr(op=op, lhs=lhs, rhs=rhs)
return result
def visit_BoolOp(self, node: ast.BoolOp) -> gt_ir.BinOpExpr:
op = self.visit(node.op)
lhs = gt_ir.utils.make_expr(self.visit(node.values[0]))
args = [lhs]
assert len(node.values) >= 2
rhs = gt_ir.utils.make_expr(self.visit(node.values[-1]))
args.append(rhs)
for i in range(len(node.values) - 2, 0, -1):
lhs = gt_ir.utils.make_expr(self.visit(node.values[i]))
rhs = gt_ir.BinOpExpr(op=op, lhs=lhs, rhs=rhs)
args.append(lhs)
result = gt_ir.BinOpExpr(op=op, lhs=lhs, rhs=rhs)
return result
def visit_Compare(self, node: ast.Compare) -> gt_ir.BinOpExpr:
lhs = gt_ir.utils.make_expr(self.visit(node.left))
args = [lhs]
assert len(node.comparators) >= 1
op = self.visit(node.ops[-1])
rhs = gt_ir.utils.make_expr(self.visit(node.comparators[-1]))
args.append(rhs)
for i in range(len(node.comparators) - 2, -1, -1):
lhs = gt_ir.utils.make_expr(self.visit(node.values[i]))
rhs = gt_ir.BinOpExpr(op=op, lhs=lhs, rhs=rhs)
op = self.visit(node.ops[i])
args.append(lhs)
result = gt_ir.BinOpExpr(op=op, lhs=lhs, rhs=rhs)
return result
def visit_AugAssign(self, node: gt_ir.AugAssign) -> SIR.AssignmentExpr:
# ignore types due to attribclass problem
bin_op = gt_ir.BinOpExpr(lhs=node.target, op=node.op,
rhs=node.value) # type: ignore
assign = gt_ir.Assign(target=node.target, value=bin_op) # type: ignore
return self.visit(assign)
def visit_AugAssign(self, node: gt_ir.AugAssign):
bin_op = gt_ir.BinOpExpr(lhs=node.target, op=node.op, rhs=node.value)
assign = gt_ir.Assign(target=node.target, value=bin_op)
return self.visit(assign)