Qualcomm AI Engine Direct - Decompose integer floor_divide by a Python scalar

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):

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(