Skip to content

Commit a79f21e

Browse files
gokulkrishna98gokulkrishna98
committed
_aten_to_core: implement aten::atan2
Adds a converter for aten.atan2.default. coreai has no native atan2 op, so the conversion decomposes it into atan(y/x) with per-quadrant correction: - x > 0: atan(y/x) - x < 0, y ≥ 0: atan(y/x) + π - x < 0, y < 0: atan(y/x) − π - x = 0, y > 0: π/2 - x = 0, y < 0: −π/2 - x = 0, y = 0: 0 Division by zero when x=0 is guarded with broadcasting_where before the divide, then the x=0 result is selected in a final where at the end. Adds numerical tests (shapes 1D/2D/3D, float32/float16, static/dynamic, axis-aligned edge cases) and IR FileCheck tests (static, dynamic, 1D).
1 parent ea728d6 commit a79f21e

3 files changed

Lines changed: 216 additions & 7 deletions

File tree

‎coreai_torch/_aten_to_core.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1522,6 +1522,46 @@ def replace_argmax(values_map: dict[str, Value], node: fx.Node, loc: Location) -
15221522
return result if keepdim else coreai.shrink_dims(result, [dim])
15231523

15241524

1525+
def replace_atan2(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
1526+
"""atan2(y, x) via atan(y/x) with quadrant correction.
1527+
1528+
atan2 is undefined for (y=0, x=0); follows the convention atan2(0, 0) = 0.
1529+
"""
1530+
y, x = _get_operands(values_map, node, [0, 1])
1531+
ele_type = x.type.element_type
1532+
1533+
zero = coreai.constant(0.0, dtype=ele_type)
1534+
pi = coreai.constant(np.pi, dtype=ele_type)
1535+
half_pi = coreai.constant(np.pi / 2.0, dtype=ele_type)
1536+
neg_half_pi = coreai.constant(-np.pi / 2.0, dtype=ele_type)
1537+
1538+
# Avoid division by zero when x = 0 by substituting x = 1 for the ratio.
1539+
x_is_zero = coreai.broadcasting_equal(x, zero)
1540+
x_safe = coreai.broadcasting_where(
1541+
x_is_zero, coreai.constant(1.0, dtype=ele_type), x
1542+
)
1543+
base = coreai.atan(coreai.broadcasting_divide(y, x_safe))
1544+
1545+
# Quadrant correction: x < 0 shifts the result by ±π.
1546+
x_neg = coreai.broadcasting_greater(zero, x)
1547+
y_neg = coreai.broadcasting_greater(zero, y)
1548+
y_pos = coreai.broadcasting_greater(y, zero)
1549+
correction = coreai.broadcasting_where(
1550+
y_neg,
1551+
coreai.broadcasting_sub(base, pi),
1552+
coreai.broadcasting_add(base, pi),
1553+
)
1554+
nonzero_result = coreai.broadcasting_where(x_neg, correction, base)
1555+
1556+
# x = 0: result is π/2, −π/2, or 0 based on sign of y.
1557+
zero_result = coreai.broadcasting_where(
1558+
y_pos,
1559+
half_pi,
1560+
coreai.broadcasting_where(y_neg, neg_half_pi, zero),
1561+
)
1562+
return coreai.broadcasting_where(x_is_zero, zero_result, nonzero_result)
1563+
1564+
15251565
def replace_gather(values_map: dict[str, Value], node: fx.Node, loc: Location) -> Value:
15261566
"""Converts aten.gather to coreai.gather_along_axis."""
15271567
x, index = _get_operands(values_map, node, [0, 2])
@@ -3440,6 +3480,7 @@ def sdpa_maskless(q: Value, k: Value, v: Value) -> Value:
34403480
"asin.default": replace_unary_ops,
34413481
"asinh.default": replace_unary_ops,
34423482
"atan.default": replace_unary_ops,
3483+
"atan2.default": replace_atan2,
34433484
"atanh.default": replace_unary_ops,
34443485
"_adaptive_avg_pool2d.default": replace_adaptive_avg_pool2d,
34453486
"_unsafe_view.default": replace_view,

‎tests/ops/test_ops.py‎

Lines changed: 58 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -595,13 +595,64 @@ def forward(self, x: Tensor) -> Tensor:
595595
await validate_numerical_output(model=model, x=x, dynamic_shapes=dynamic_shapes)
596596

597597

598-
@pytest.mark.parametrize(
599-
"x",
600-
[
601-
torch.rand(2, 3, 8, 8),
602-
torch.rand(2, 3, 8, 8, dtype=torch.float16), # fp16
603-
],
604-
)
598+
class TestAtan2:
599+
"""Tests for torch.atan2(y, x) — angle from the positive x-axis to the point (x, y)."""
600+
601+
class Atan2Model(nn.Module):
602+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
603+
return torch.atan2(y, x)
604+
605+
@pytest.mark.parametrize("dynamic", [False, True])
606+
@pytest.mark.parametrize(
607+
"shape",
608+
[
609+
(4,),
610+
(3, 4),
611+
(2, 3, 4),
612+
],
613+
)
614+
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16])
615+
async def test_basic(
616+
self, shape: tuple[int, ...], dtype: torch.dtype, dynamic: bool
617+
) -> None:
618+
model = self.Atan2Model().eval()
619+
y = torch.randn(shape, dtype=dtype)
620+
x = torch.randn(shape, dtype=dtype)
621+
dynamic_shapes = (
622+
{"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)} if dynamic else None
623+
)
624+
await validate_numerical_output(
625+
model=model, y=y, x=x, dynamic_shapes=dynamic_shapes
626+
)
627+
628+
async def test_x_zero(self) -> None:
629+
"""x = 0 should yield ±π/2 depending on sign of y."""
630+
model = self.Atan2Model().eval()
631+
y = torch.tensor([1.0, -1.0, 2.0, -2.0])
632+
x = torch.zeros(4)
633+
await validate_numerical_output(model=model, y=y, x=x)
634+
635+
async def test_y_zero(self) -> None:
636+
"""y = 0 with x > 0 → 0, x < 0 → π."""
637+
model = self.Atan2Model().eval()
638+
y = torch.zeros(4)
639+
x = torch.tensor([1.0, -1.0, 2.0, -2.0])
640+
await validate_numerical_output(model=model, y=y, x=x)
641+
642+
async def test_all_quadrants(self) -> None:
643+
"""Cover all four quadrants and axes."""
644+
model = self.Atan2Model().eval()
645+
y = torch.tensor([1.0, 1.0, -1.0, -1.0, 0.0, 0.0, 1.0, -1.0])
646+
x = torch.tensor([1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 0.0, 0.0])
647+
await validate_numerical_output(model=model, y=y, x=x)
648+
649+
async def test_broadcast_shapes(self) -> None:
650+
model = self.Atan2Model().eval()
651+
y = torch.randn(3, 4)
652+
x = torch.randn(4)
653+
await validate_numerical_output(model=model, y=y, x=x)
654+
655+
605656
@pytest.mark.parametrize(
606657
"dynamic_dims", [tuple(), (0,), (2,), (3,), (0, 2), (0, 3), (0, 2, 3)]
607658
)

‎tests/ops/test_ops_ir.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1132,6 +1132,123 @@ def forward(self, x: Tensor) -> Tensor:
11321132
)
11331133

11341134

1135+
class TestAtan2IR:
1136+
def test_static(self) -> None:
1137+
class Atan2Model(nn.Module):
1138+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1139+
return torch.atan2(y, x)
1140+
1141+
ir = get_ir(Atan2Model().eval(), y=torch.rand(2, 3), x=torch.rand(2, 3))
1142+
filecheck_pattern(
1143+
ir,
1144+
check_file="""
1145+
// CHECK-LABEL: module {
1146+
// CHECK-NEXT: coreai.graph @main(%[[ARG0:.*]]: tensor<2x3xf32> {coreai.name = "y"}, %[[ARG1:.*]]: tensor<2x3xf32> {coreai.name = "x"}) -> (tensor<2x3xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
1147+
// CHECK-NEXT: %[[C1:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1148+
// CHECK-NEXT: %[[C0:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1149+
// CHECK-NEXT: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
1150+
// CHECK-NEXT: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
1151+
// CHECK-NEXT: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
1152+
// CHECK-NEXT: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[ARG1]], %[[C0]] : (tensor<2x3xf32>, tensor<f32>) -> tensor<2x3xi1>
1153+
// CHECK-NEXT: %[[X_SAFE:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[C1]], %[[ARG1]] : (tensor<2x3xi1>, tensor<f32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1154+
// CHECK-NEXT: %[[RATIO:.*]] = coreai.decomposable.broadcasting_divide %[[ARG0]], %[[X_SAFE]] : (tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1155+
// CHECK-NEXT: %[[BASE:.*]] = coreai.atan %[[RATIO]] : tensor<2x3xf32> -> tensor<2x3xf32>
1156+
// CHECK-NEXT: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG1]] : (tensor<f32>, tensor<2x3xf32>) -> tensor<2x3xi1>
1157+
// CHECK-NEXT: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG0]] : (tensor<f32>, tensor<2x3xf32>) -> tensor<2x3xi1>
1158+
// CHECK-NEXT: %[[Y_POS:.*]] = coreai.decomposable.broadcasting_greater %[[ARG0]], %[[C0]] : (tensor<2x3xf32>, tensor<f32>) -> tensor<2x3xi1>
1159+
// CHECK-NEXT: %[[BASE_MINUS_PI:.*]] = coreai.decomposable.broadcasting_sub %[[BASE]], %[[PI]] : (tensor<2x3xf32>, tensor<f32>) -> tensor<2x3xf32>
1160+
// CHECK-NEXT: %[[BASE_PLUS_PI:.*]] = coreai.decomposable.broadcasting_add %[[BASE]], %[[PI]] : (tensor<2x3xf32>, tensor<f32>) -> tensor<2x3xf32>
1161+
// CHECK-NEXT: %[[CORRECTION:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[BASE_MINUS_PI]], %[[BASE_PLUS_PI]] : (tensor<2x3xi1>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1162+
// CHECK-NEXT: %[[NONZERO:.*]] = coreai.decomposable.broadcasting_where %[[X_NEG]], %[[CORRECTION]], %[[BASE]] : (tensor<2x3xi1>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1163+
// CHECK-NEXT: %[[ZERO_NEG:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[NHPI]], %[[C0]] : (tensor<2x3xi1>, tensor<f32>, tensor<f32>) -> tensor<2x3xf32>
1164+
// CHECK-NEXT: %[[ZERO_RES:.*]] = coreai.decomposable.broadcasting_where %[[Y_POS]], %[[HPI]], %[[ZERO_NEG]] : (tensor<2x3xi1>, tensor<f32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1165+
// CHECK-NEXT: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[ZERO_RES]], %[[NONZERO]] : (tensor<2x3xi1>, tensor<2x3xf32>, tensor<2x3xf32>) -> tensor<2x3xf32>
1166+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<2x3xf32>
1167+
// CHECK-NEXT: }
1168+
// CHECK-NEXT: }
1169+
""",
1170+
)
1171+
1172+
def test_dynamic(self) -> None:
1173+
class Atan2Model(nn.Module):
1174+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1175+
return torch.atan2(y, x)
1176+
1177+
y = torch.rand(2, 3)
1178+
x = torch.rand(2, 3)
1179+
ir = get_ir(
1180+
Atan2Model().eval(),
1181+
y=y,
1182+
x=x,
1183+
dynamic_shapes={"y": _all_dims_dynamic(y), "x": _all_dims_dynamic(x)},
1184+
)
1185+
filecheck_pattern(
1186+
ir,
1187+
check_file="""
1188+
// CHECK-LABEL: module {
1189+
// CHECK-NEXT: coreai.graph @main(%[[ARG0:.*]]: tensor<?x?xf32> {coreai.name = "y"}, %[[ARG1:.*]]: tensor<?x?xf32> {coreai.name = "x"}) -> (tensor<?x?xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
1190+
// CHECK-NEXT: %[[C1:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1191+
// CHECK-NEXT: %[[C0:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1192+
// CHECK-NEXT: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
1193+
// CHECK-NEXT: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
1194+
// CHECK-NEXT: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
1195+
// CHECK-NEXT: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[ARG1]], %[[C0]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
1196+
// CHECK-NEXT: %[[X_SAFE:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[C1]], %[[ARG1]] : (tensor<?x?xi1>, tensor<f32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1197+
// CHECK-NEXT: %[[RATIO:.*]] = coreai.decomposable.broadcasting_divide %[[ARG0]], %[[X_SAFE]] : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1198+
// CHECK-NEXT: %[[BASE:.*]] = coreai.atan %[[RATIO]] : tensor<?x?xf32> -> tensor<?x?xf32>
1199+
// CHECK-NEXT: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG1]] : (tensor<f32>, tensor<?x?xf32>) -> tensor<?x?xi1>
1200+
// CHECK-NEXT: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG0]] : (tensor<f32>, tensor<?x?xf32>) -> tensor<?x?xi1>
1201+
// CHECK-NEXT: %[[Y_POS:.*]] = coreai.decomposable.broadcasting_greater %[[ARG0]], %[[C0]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xi1>
1202+
// CHECK-NEXT: %[[BASE_MINUS_PI:.*]] = coreai.decomposable.broadcasting_sub %[[BASE]], %[[PI]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
1203+
// CHECK-NEXT: %[[BASE_PLUS_PI:.*]] = coreai.decomposable.broadcasting_add %[[BASE]], %[[PI]] : (tensor<?x?xf32>, tensor<f32>) -> tensor<?x?xf32>
1204+
// CHECK-NEXT: %[[CORRECTION:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[BASE_MINUS_PI]], %[[BASE_PLUS_PI]] : (tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1205+
// CHECK-NEXT: %[[NONZERO:.*]] = coreai.decomposable.broadcasting_where %[[X_NEG]], %[[CORRECTION]], %[[BASE]] : (tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1206+
// CHECK-NEXT: %[[ZERO_NEG:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[NHPI]], %[[C0]] : (tensor<?x?xi1>, tensor<f32>, tensor<f32>) -> tensor<?x?xf32>
1207+
// CHECK-NEXT: %[[ZERO_RES:.*]] = coreai.decomposable.broadcasting_where %[[Y_POS]], %[[HPI]], %[[ZERO_NEG]] : (tensor<?x?xi1>, tensor<f32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1208+
// CHECK-NEXT: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[ZERO_RES]], %[[NONZERO]] : (tensor<?x?xi1>, tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
1209+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<?x?xf32>
1210+
// CHECK-NEXT: }
1211+
// CHECK-NEXT: }
1212+
""",
1213+
)
1214+
1215+
def test_1d(self) -> None:
1216+
class Atan2Model(nn.Module):
1217+
def forward(self, y: Tensor, x: Tensor) -> Tensor:
1218+
return torch.atan2(y, x)
1219+
1220+
ir = get_ir(Atan2Model().eval(), y=torch.rand(4), x=torch.rand(4))
1221+
filecheck_pattern(
1222+
ir,
1223+
check_file="""
1224+
// CHECK-LABEL: module {
1225+
// CHECK-NEXT: coreai.graph @main(%[[ARG0:.*]]: tensor<4xf32> {coreai.name = "y"}, %[[ARG1:.*]]: tensor<4xf32> {coreai.name = "x"}) -> (tensor<4xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
1226+
// CHECK-NEXT: %[[C1:.*]] = coreai.constant dense<1.000000e+00> : tensor<f32>
1227+
// CHECK-NEXT: %[[C0:.*]] = coreai.constant dense<0.000000e+00> : tensor<f32>
1228+
// CHECK-NEXT: %[[PI:.*]] = coreai.constant dense<3.14159274> : tensor<f32>
1229+
// CHECK-NEXT: %[[HPI:.*]] = coreai.constant dense<1.57079637> : tensor<f32>
1230+
// CHECK-NEXT: %[[NHPI:.*]] = coreai.constant dense<-1.57079637> : tensor<f32>
1231+
// CHECK-NEXT: %[[X_IS_ZERO:.*]] = coreai.decomposable.broadcasting_equal %[[ARG1]], %[[C0]] : (tensor<4xf32>, tensor<f32>) -> tensor<4xi1>
1232+
// CHECK-NEXT: %[[X_SAFE:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[C1]], %[[ARG1]] : (tensor<4xi1>, tensor<f32>, tensor<4xf32>) -> tensor<4xf32>
1233+
// CHECK-NEXT: %[[RATIO:.*]] = coreai.decomposable.broadcasting_divide %[[ARG0]], %[[X_SAFE]] : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
1234+
// CHECK-NEXT: %[[BASE:.*]] = coreai.atan %[[RATIO]] : tensor<4xf32> -> tensor<4xf32>
1235+
// CHECK-NEXT: %[[X_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG1]] : (tensor<f32>, tensor<4xf32>) -> tensor<4xi1>
1236+
// CHECK-NEXT: %[[Y_NEG:.*]] = coreai.decomposable.broadcasting_greater %[[C0]], %[[ARG0]] : (tensor<f32>, tensor<4xf32>) -> tensor<4xi1>
1237+
// CHECK-NEXT: %[[Y_POS:.*]] = coreai.decomposable.broadcasting_greater %[[ARG0]], %[[C0]] : (tensor<4xf32>, tensor<f32>) -> tensor<4xi1>
1238+
// CHECK-NEXT: %[[BASE_MINUS_PI:.*]] = coreai.decomposable.broadcasting_sub %[[BASE]], %[[PI]] : (tensor<4xf32>, tensor<f32>) -> tensor<4xf32>
1239+
// CHECK-NEXT: %[[BASE_PLUS_PI:.*]] = coreai.decomposable.broadcasting_add %[[BASE]], %[[PI]] : (tensor<4xf32>, tensor<f32>) -> tensor<4xf32>
1240+
// CHECK-NEXT: %[[CORRECTION:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[BASE_MINUS_PI]], %[[BASE_PLUS_PI]] : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
1241+
// CHECK-NEXT: %[[NONZERO:.*]] = coreai.decomposable.broadcasting_where %[[X_NEG]], %[[CORRECTION]], %[[BASE]] : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
1242+
// CHECK-NEXT: %[[ZERO_NEG:.*]] = coreai.decomposable.broadcasting_where %[[Y_NEG]], %[[NHPI]], %[[C0]] : (tensor<4xi1>, tensor<f32>, tensor<f32>) -> tensor<4xf32>
1243+
// CHECK-NEXT: %[[ZERO_RES:.*]] = coreai.decomposable.broadcasting_where %[[Y_POS]], %[[HPI]], %[[ZERO_NEG]] : (tensor<4xi1>, tensor<f32>, tensor<4xf32>) -> tensor<4xf32>
1244+
// CHECK-NEXT: %[[RESULT:.*]] = coreai.decomposable.broadcasting_where %[[X_IS_ZERO]], %[[ZERO_RES]], %[[NONZERO]] : (tensor<4xi1>, tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
1245+
// CHECK-NEXT: coreai.output %[[RESULT]] : tensor<4xf32>
1246+
// CHECK-NEXT: }
1247+
// CHECK-NEXT: }
1248+
""",
1249+
)
1250+
1251+
11351252
class TestAvgPool2dIR:
11361253
def test_static(self) -> None:
11371254
class AvgPool2dModel(nn.Module):

0 commit comments

Comments
 (0)