diff --git a/backends/qualcomm/_passes/decompose_floor_divide.py b/backends/qualcomm/_passes/decompose_floor_divide.py index bf7db3d5b95..d059a486990 100644 --- a/backends/qualcomm/_passes/decompose_floor_divide.py +++ b/backends/qualcomm/_passes/decompose_floor_divide.py @@ -12,14 +12,14 @@ class FloorDivide(torch.nn.Module): - def __init__(self): + def __init__(self, dtype: torch.dtype): super().__init__() + self.dtype = dtype def forward(self, x, y): - dtype = x.dtype result = torch.div(x, y) result = torch.floor(result) - return result.to(dtype) + return result.to(self.dtype) class DecomposeFloorDivide(ExportPass): @@ -36,14 +36,23 @@ def __init__(self) -> None: def call(self, graph_module: torch.fx.GraphModule) -> PassResult: graph = graph_module.graph for node in graph.nodes: - model = FloorDivide() if ( torch.ops.aten.floor_divide.default == node.target and not torch.is_floating_point(node.meta["val"]) ): + # The divisor may be a Python scalar (e.g. `indices // num_classes`): + # trace it as a 0-dim tensor, then remap it back to the scalar. + model = FloorDivide(node.meta["val"].dtype) decomposed_module = torch.export.export( model, - (node.args[0].meta["val"], node.args[1].meta["val"]), + tuple( + ( + arg.meta["val"] + if isinstance(arg, torch.fx.Node) + else torch.tensor(arg) + ) + for arg in node.args[:2] + ), strict=True, ).module() with graph.inserting_before(node): diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..bc043c4368f 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -2013,6 +2013,10 @@ class _FloatDiv(torch.nn.Module): def forward(self, x, y): return torch.floor_divide(x, y) + class _IntDivScalar(torch.nn.Module): + def forward(self, x): + return x // 3 + @staticmethod @unpack_pass_fixtures def test( @@ -2041,6 +2045,18 @@ def test( assertions.assert_no_target(gm, torch.ops.aten.floor_divide.default) assertions.assert_target_count(gm, torch.ops.aten.floor.default, 1) + # Python scalar divisor, as in `indices // num_classes` + with subtests.test(msg="int_scalar"): + gm = pass_pipeline.lower_export_gm( + module=DecomposeFloorDivide._IntDivScalar(), + sample_input=(torch.randint(-10, 10, (3, 4), dtype=torch.int32),), + target_pass=target_pass, + backend_type=backend_type, + quantizer=quantizer, + ) + assertions.assert_no_target(gm, torch.ops.aten.floor_divide.default) + assertions.assert_target_count(gm, torch.ops.aten.floor.default, 1) + # Float inputs: pass does not fire — floor_divide preserved as-is with subtests.test(msg="float"): gm = pass_pipeline.lower_export_gm(