add ut for int8 scaled_embeding_bag pattern_match

Author: shiyang-weng

test/quantization/pt2e/test_x86inductor_fusion.py

@@ -3047,42 +3047,60 @@ def test_fp8_q_attention_block(self):
                 annotate_matmul=annotate_matmul, is_fp8=True
             )
 
-    @skipIfNoDynamoSupport
-    @skipIfNoONEDNN
-    @skipIfNoFloat8Support
-    @unittest.skipIf(
-        "CPU" not in torch._C._dispatch_dump("torchao::_scaled_embedding_bag"),
-        reason="cpp kernels not built",
-    )
-    def test_fp8_scaled_embedding_bag(self):
-        dtype = torch.float8_e4m3fn
-
+    def _test_scaled_embedding_bag_helper(self, dtype, with_quant=False):
         class FP8QDQEmbeddingBag(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.weight_scale = 2.0
+                self.output_scale = 3.0
+
+            def _dq(self, weight):
+                if dtype == torch.float8_e4m3fn:
+                    res = torch.ops.torchao.dequantize_affine_float8_non_decomposed.default(
+                        tensor=weight.data,
+                        scale=torch.tensor([self.weight_scale]),
+                        output_dtype=torch.float,
+                    )
+                else:
+                    res = torch.ops.quantized_decomposed.dequantize_per_tensor.default(
+                        weight.data,
+                        self.weight_scale,
+                        0,
+                        -128, 127, torch.int8,
+                    )
+                return res
+
+            def _q(self, x):
+                if dtype == torch.float8_e4m3fn:
+                    qx = torch.ops.torchao.quantize_affine_float8_non_decomposed.default(
+                        tensor=x,
+                        scale=torch.tensor([self.output_scale]),
+                        float8_dtype=dtype,
+                    )
+                else:
+                    qx = torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                        x, self.output_scale, 0, -128, 127, torch.int8
+                    )
+                return qx
 
             def forward(
                 self,
                 weight,
                 input,
                 offsets=None,
             ):
-                weight = (
-                    torch.ops.torchao.dequantize_affine_float8_non_decomposed.default(
-                        tensor=weight.data,
-                        scale=torch.tensor([self.weight_scale]),
-                        output_dtype=torch.float,
-                    )
-                )
+                weight = self._dq(weight)
 
-                return torch.nn.functional.embedding_bag(
+                res = torch.nn.functional.embedding_bag(
                     input,
                     weight,
                     offsets,
                     mode="sum",
                     include_last_offset=True,
                 )
+                if with_quant:
+                    res = self._q(res)
+                return res
 
         EMBEDINGBAG_MULTIHOT_SIZES = [1, 2, 3, 10]
         EMBEDINGBAG_BAG_SIZES = [1, 2, 128, 1024]
@@ -3109,8 +3127,11 @@ def forward(
                 )
 
                 def matcher_check_fn():
+                    counter_name = "scaled_embedding_bag"
+                    if with_quant:
+                        counter_name += "_with_quant"
                     self.assertEqual(
-                        counters["inductor"]["scaled_embedding_bag_matcher_count"], 1
+                        counters["inductor"][f"{counter_name}_matcher_count"], 1
                     )
 
                 self._test_common(
@@ -3120,6 +3141,38 @@ def matcher_check_fn():
                 )
 
 
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNN
+    @skipIfNoFloat8Support
+    @unittest.skipIf(
+        "CPU" not in torch._C._dispatch_dump("torchao::_scaled_embedding_bag"),
+        reason="cpp kernels not built",
+    )
+    def test_fp8_scaled_embedding_bag(self):
+        self._test_scaled_embedding_bag_helper(torch.float8_e4m3fn)
+
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNN
+    @skipIfNoFloat8Support
+    @unittest.skipIf(
+        "CPU" not in torch._C._dispatch_dump("torchao::_scaled_embedding_bag"),
+        reason="cpp kernels not built",
+    )
+    def test_int8_scaled_embedding_bag(self):
+        self._test_scaled_embedding_bag_helper(torch.int8)
+
+
+    @skipIfNoDynamoSupport
+    @skipIfNoONEDNN
+    @skipIfNoFloat8Support
+    @unittest.skipIf(
+        "CPU" not in torch._C._dispatch_dump("torchao::_scaled_embedding_bag"),
+        reason="cpp kernels not built",
+    )
+    def test_int8_scaled_embedding_bag_with_quant(self):
+        self._test_scaled_embedding_bag_helper(torch.int8, True)
+
+
 instantiate_parametrized_tests(TestPatternMatcher)
 if __name__ == "__main__":
     if IS_LINUX and HAS_CPU and torch.backends.mkldnn.is_available():

torchao/quantization/pt2e/inductor_passes/x86.py

@@ -2893,20 +2893,17 @@ def scaled_embedding_bag(match: Match, *args, **kwargs):
             kwargs["mode"],
             kwargs["include_last_offset"],
         )
-        # only support fp32 and int8 output on kernel
-        # next step to support more output_type
         output_type = torch.float
         o_scale = 1.0
         if "o_dtype" in kwargs:
-            output_type = torch.int8
+            output_type = kwargs["o_dtype"]
             o_scale = kwargs["o_inv_scale"]
 
         graph = match.graph
         with graph.inserting_before(getitem_node):
-            # scale type is float on int8 q/dq
-            # Not support float scale yet on scaled_embedding_bag
+            # float scale not supported on scaled_embedding_bag
             # convert scale from float into tensor
-            if output_type == torch.int8:
+            if type(w_scale) is float:
                 w_scale = graph.call_function(
                     torch.ops.aten.full.default,
                     args=([1], w_scale),
@@ -2927,7 +2924,7 @@ def scaled_embedding_bag(match: Match, *args, **kwargs):
                 torch.ops.torchao._scaled_embedding_bag.default, args=new_args
             )
 
-            # remove quant node
+            # Erase quant pattern
             if output_type == torch.int8:
                 quant_node.replace_all_uses_with(getitem_node)
                 getitem_node.meta.update(quant_node.meta)
@@ -2942,8 +2939,11 @@ def scaled_embedding_bag(match: Match, *args, **kwargs):
             # Erase the dequant pattern
             graph.erase_node(dequant_node)
 
-        counters["inductor"]["scaled_embedding_bag_matcher_count"] += 1
-        counters["inductor"]["scaled_embedding_bag_matcher_nodes"] += len(match.nodes)
+        counter_name = "scaled_embedding_bag"
+        if "o_dtype" in kwargs:
+            counter_name += "_with_quant"
+        counters["inductor"][f"{counter_name}_matcher_count"] += 1
+        counters["inductor"][f"{counter_name}_matcher_nodes"] += len(match.nodes)
 
 
 def _generate_scaled_embedding_bag_patterns(dq_pattern):