Skip to content

Commit c0df044

Browse files
agrima1304agrima1304
authored
Arm backend: Add support for floor_divide.default (#14933)
Converts floor_divide.default to div.Tensor_mode with rounding_mode = "floor" Signed-off-by: Agrima Khare <[email protected]>
1 parent 9d68039 commit c0df044

File tree

5 files changed

+235
-0
lines changed

5 files changed

+235
-0
lines changed

‎backends/arm/_passes/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
from .decompose_elu_pass import DecomposeEluPass # noqa
4343
from .decompose_embedding_pass import DecomposeEmbeddingPass # noqa # noqa
4444
from .decompose_expm1_pass import DecomposeExpm1Pass # noqa
45+
from .decompose_floor_divide_pass import DecomposeFloorDividePass # noqa
4546
from .decompose_gelu_pass import DecomposeGeluPass # noqa
4647
from .decompose_glu_pass import DecomposeGluPass # noqa
4748
from .decompose_grouped_conv import DecomposeGroupedConv # noqa

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
DecomposeEluPass,
5050
DecomposeEmbeddingPass,
5151
DecomposeExpm1Pass,
52+
DecomposeFloorDividePass,
5253
DecomposeGeluPass,
5354
DecomposeGluPass,
5455
DecomposeGroupedConv,
@@ -245,6 +246,8 @@ def _tosa_FP_pipeline(
245246
self.add_pass(CastBoolToInt8Pass())
246247
self.add_pass(DecomposeSinhPass())
247248
self.add_pass(DecomposeSignPass())
249+
self.add_pass(DecomposeFloorDividePass())
250+
self.add_pass(DecomposeDivTensorModePass())
248251
self.add_pass(ReplaceScalarWithTensorByProfilePass())
249252
self.add_pass(DecomposeRemainderPass())
250253
self.add_pass(DecomposeDivTensorModePass())
@@ -339,6 +342,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
339342
self.add_pass(DecomposeAddmmPass())
340343
self.add_pass(ReplaceScalarWithTensorByProfilePass())
341344
self.add_pass(DecomposeRemainderPass())
345+
self.add_pass(DecomposeFloorDividePass())
342346
self.add_pass(DecomposeDivTensorModePass())
343347
self.add_pass(DecomposeAddSubAlphaPass())
344348
self.add_pass(ScalarsToAttributePass())

‎backends/arm/_passes/decompose_floor_divide_pass.py‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
# Copyright 2025 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from typing import Set, Type
7+
8+
import torch
9+
from executorch.backends.arm._passes import ArmPass
10+
from executorch.backends.arm._passes.decompose_div_tensor_mode import (
11+
DecomposeDivTensorModePass,
12+
)
13+
from executorch.exir.dialects._ops import ops as exir_ops
14+
from executorch.exir.pass_base import ExportPass
15+
16+
edge_floor_divide_ops = (exir_ops.edge.aten.floor_divide.default,)
17+
aten_floor_divide_ops = (torch.ops.aten.floor_divide.default,)
18+
19+
20+
def get_floor_divide_decomposition(op) -> tuple:
21+
"""
22+
Returns the decomposition of the given aten.floor_div operation into
23+
its equivalent TOSA-supported operations
24+
25+
This handles both edge dialect ops and core PyTorch ops. The decomposition strategy
26+
is:
27+
floor_div(x, y) → div_tensor_mode(x, y, rounding_mode="floor")
28+
29+
Returns:
30+
A tuple (div_op,) corresponding to the appropriate operator overload for the input op.
31+
32+
Raises:
33+
RuntimeError: If the provided operator is not a supported floor_divide variant.
34+
"""
35+
36+
if op in edge_floor_divide_ops:
37+
return (
38+
exir_ops.edge.aten.div.Tensor_mode,
39+
exir_ops.edge.aten.full_like.default,
40+
)
41+
if op in aten_floor_divide_ops:
42+
return (
43+
torch.ops.aten.div.Tensor_mode,
44+
torch.ops.aten.full_like.default,
45+
)
46+
47+
raise RuntimeError(f"Can't get floor_div decomposition for op {op}")
48+
49+
50+
class DecomposeFloorDividePass(ArmPass):
51+
"""
52+
Decomposes aten.floor_divide into aten.div.Tensor_mode with rounding_mode="floor".
53+
"""
54+
55+
_passes_required_after: Set[Type[ExportPass]] = {DecomposeDivTensorModePass}
56+
57+
def call_operator(self, op, args, kwargs, meta):
58+
if op not in (edge_floor_divide_ops + aten_floor_divide_ops):
59+
return super().call_operator(op, args, kwargs, meta, updated=False)
60+
61+
(div_op, full_op) = get_floor_divide_decomposition(op)
62+
63+
input = args[0]
64+
other = args[1]
65+
66+
if isinstance(other, int):
67+
other = super().call_operator(
68+
full_op, (input, other), {}, meta, updated=False
69+
)
70+
71+
div_node = super().call_operator(
72+
div_op, (input, other), {"rounding_mode": "floor"}, meta, updated=True
73+
)
74+
75+
return div_node

‎backends/arm/operator_support/tosa_profile_supported_op_lists.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -235,6 +235,7 @@
235235
exir_ops.edge.aten.acos.default,
236236
exir_ops.edge.aten.elu.default,
237237
exir_ops.edge.aten.copy.default,
238+
exir_ops.edge.aten.floor_divide.default,
238239
}
239240

240241

‎backends/arm/test/ops/test_floor_div.py‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
# All rights reserved.
3+
# Copyright 2024-2025 Arm Limited and/or its affiliates.
4+
#
5+
# This source code is licensed under the BSD-style license found in the
6+
# LICENSE file in the root directory of this source tree.
7+
8+
from typing import Tuple, Union
9+
10+
import torch
11+
from executorch.backends.arm.test import common
12+
13+
from executorch.backends.arm.test.tester.test_pipeline import (
14+
EthosU55PipelineINT,
15+
EthosU85PipelineINT,
16+
TosaPipelineFP,
17+
TosaPipelineINT,
18+
VgfPipeline,
19+
)
20+
21+
test_data_suite = {
22+
# (test_name, input, other)
23+
"op_floor_div_rank1_ones": lambda: (
24+
torch.ones(5),
25+
torch.ones(5),
26+
),
27+
"op_floor_div_rank1_rand": lambda: (
28+
torch.rand(5) * 5,
29+
torch.rand(5) * 5,
30+
),
31+
"op_floor_div_rank4_negative_ones": lambda: (
32+
(-1) * torch.ones(5, 10, 25, 20),
33+
torch.ones(5, 10, 25, 20),
34+
),
35+
"op_floor_div_rank4_ones_div_negative": lambda: (
36+
torch.ones(5, 10, 25, 20),
37+
(-1) * torch.ones(5, 10, 25, 20),
38+
),
39+
"op_floor_div_rank4_randn_mutltiple_broadcasts": lambda: (
40+
torch.randn(1, 4, 4, 1),
41+
torch.randn(1, 1, 4, 4),
42+
),
43+
"op_floor_div_rank4_randn_scalar": lambda: (
44+
torch.randn(1, 4, 4, 1),
45+
2,
46+
),
47+
"op_floor_div_rank4_large_rand": lambda: (
48+
200 * torch.rand(5, 10, 25, 20),
49+
torch.rand(5, 10, 25, 20),
50+
),
51+
}
52+
53+
54+
class FloorDivide(torch.nn.Module):
55+
aten_op = "torch.ops.aten.floor_divide.default"
56+
aten_ops_int = ["aten.mul.Tensor", "aten.reciprocal.default", "aten.floor.default"]
57+
exir_op = "executorch_exir_dialects_edge__ops_aten_div_Tensor_mode"
58+
exir_ops_int = [
59+
"executorch_exir_dialects_edge__ops_aten_reciprocal_default",
60+
"executorch_exir_dialects_edge__ops_aten_mul_Tensor",
61+
"executorch_exir_dialects_edge__ops_aten_floor_default",
62+
]
63+
64+
def forward(
65+
self,
66+
input_: Union[torch.Tensor, torch.types.Number],
67+
other_: Union[torch.Tensor, torch.types.Number],
68+
):
69+
return torch.floor_divide(input=input_, other=other_)
70+
71+
72+
input_t1 = Tuple[torch.Tensor, Union[torch.Tensor, int]]
73+
74+
75+
@common.parametrize("test_data", test_data_suite)
76+
def test_floor_divide_tosa_FP(test_data: input_t1):
77+
pipeline = TosaPipelineFP[input_t1](
78+
FloorDivide(),
79+
test_data(),
80+
FloorDivide.aten_op,
81+
FloorDivide.exir_op,
82+
use_to_edge_transform_and_lower=False,
83+
)
84+
pipeline.run()
85+
86+
87+
@common.parametrize("test_data", test_data_suite)
88+
def test_floor_divide_tosa_INT(test_data: input_t1):
89+
pipeline = TosaPipelineINT[input_t1](
90+
FloorDivide(),
91+
test_data(),
92+
aten_op=FloorDivide.aten_ops_int,
93+
exir_op=FloorDivide.exir_ops_int,
94+
use_to_edge_transform_and_lower=False,
95+
)
96+
pipeline.run()
97+
98+
99+
@common.parametrize("test_data", test_data_suite)
100+
@common.XfailIfNoCorstone300
101+
def test_floor_divide_u55_INT(test_data: input_t1):
102+
pipeline = EthosU55PipelineINT[input_t1](
103+
FloorDivide(),
104+
test_data(),
105+
aten_ops=FloorDivide.aten_ops_int,
106+
exir_ops=[],
107+
run_on_fvp=True,
108+
use_to_edge_transform_and_lower=False,
109+
)
110+
pipeline.pop_stage("check_not.exir")
111+
pipeline.pop_stage("check_count.exir")
112+
pipeline.run()
113+
114+
115+
@common.parametrize("test_data", test_data_suite)
116+
@common.XfailIfNoCorstone320
117+
def test_floor_divide_u85_INT(test_data: input_t1):
118+
pipeline = EthosU85PipelineINT[input_t1](
119+
FloorDivide(),
120+
test_data(),
121+
aten_ops=FloorDivide.aten_ops_int,
122+
exir_ops=FloorDivide.exir_ops_int,
123+
run_on_fvp=True,
124+
use_to_edge_transform_and_lower=False,
125+
)
126+
pipeline.run()
127+
128+
129+
@common.parametrize("test_data", test_data_suite)
130+
@common.SkipIfNoModelConverter
131+
def test_floor_divide_vgf_FP(test_data: input_t1):
132+
pipeline = VgfPipeline[input_t1](
133+
FloorDivide(),
134+
test_data(),
135+
FloorDivide.aten_op,
136+
FloorDivide.exir_op,
137+
tosa_version="TOSA-1.0+FP",
138+
use_to_edge_transform_and_lower=False,
139+
)
140+
pipeline.run()
141+
142+
143+
@common.parametrize("test_data", test_data_suite)
144+
@common.SkipIfNoModelConverter
145+
def test_floor_divide_vgf_INT(test_data: input_t1):
146+
pipeline = VgfPipeline[input_t1](
147+
FloorDivide(),
148+
test_data(),
149+
aten_op=FloorDivide.aten_ops_int,
150+
exir_op=FloorDivide.exir_ops_int,
151+
tosa_version="TOSA-1.0+INT",
152+
use_to_edge_transform_and_lower=False,
153+
)
154+
pipeline.run()

0 commit comments

Comments
 (0)