Extend int dtype with uint8 to skip lowering uint8-input ops in AIT s…

…plitter (#989)

Summary:
Pull Request resolved: #989

As title

Reviewed By: frank-wei, qxy11, 22quinn

Differential Revision: D53249618

fbshipit-source-id: 21643eab421eaec4376187ac6a48a3f263474d26

fx2ait/fx2ait/ait_splitter.py

@@ -85,6 +85,7 @@ def create_ait_operator_support(
         else [
             ops.OpSupports.decline_if_input_dtype(torch.int64),
             ops.OpSupports.decline_if_input_dtype(torch.int32),
+            ops.OpSupports.decline_if_input_dtype(torch.uint8),
         ]
     )
     chained_not_supported_ops += [

fx2ait/fx2ait/test/test_ait_splitter.py

@@ -165,6 +165,23 @@ def forward(self, a):
             {"_run_on_gpu_0"},
         )
 
+        # nodes w/ uint8 input should not be lowered
+        mod = acc_tracer.trace(test_mod, [x])
+        splitter = AITSplitter(
+            mod,
+            (x.to(torch.uint8).cuda(),),
+            operator_support,
+            settings,
+        )
+
+        split_results_int = splitter.generate_split_results()
+
+        self.assertTrue(len(split_results_int), 1)
+        self.assertEqual(
+            dict(split_results_int.split_module.named_children()).keys(),
+            {"_run_on_gpu_0"},
+        )
+
         # nodes w/ integer input should be lowered
         mod = acc_tracer.trace(test_mod, [x])
         settings.allow_int_inputs = True