Add fx pass to remove sdpa composite zero mask input (#62)

* init

* add copyright

* Add test

ai_edge_torch/convert/conversion.py

@@ -30,6 +30,7 @@
 from ai_edge_torch.convert.fx_passes import InjectMlirDebuginfoPass
 from ai_edge_torch.convert.fx_passes import OptimizeLayoutTransposesPass
 from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes import run_generative_passes
 from ai_edge_torch.quantize import quant_config as qcfg
 
 os.environ["EXPERIMENTAL_XLA_UNBOUNDED_DYNAMISM"] = "1"
@@ -38,6 +39,7 @@
 def _run_convert_passes(
     exported_program: ExportedProgram,
 ) -> ExportedProgram:
+  exported_program = run_generative_passes(exported_program)
   return run_passes(
       exported_program,
       [

ai_edge_torch/generative/fx_passes/__init__.py

@@ -0,0 +1,31 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+from ai_edge_torch.convert.fx_passes import CanonicalizePass
+from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes.remove_sdpa_zero_mask_pass import RemoveSDPACompositeZeroMaskPass  # NOQA
+
+
+def run_generative_passes(
+    exported_program: torch.export.ExportedProgram,
+) -> torch.export.ExportedProgram:
+  return run_passes(
+      exported_program,
+      [
+          RemoveSDPACompositeZeroMaskPass(),
+          CanonicalizePass(),
+      ],
+  )

ai_edge_torch/generative/fx_passes/remove_sdpa_zero_mask_pass.py

@@ -0,0 +1,47 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import torch
+
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassBase
+from ai_edge_torch.convert.fx_passes._pass_base import ExportedProgramPassResult  # NOQA
+
+
+class RemoveSDPACompositeZeroMaskPass(ExportedProgramPassBase):
+
+  def is_zero_tensor_node(self, node: torch.fx.Node):
+    return node.target == torch.ops.aten.zeros.default
+
+  def call(self, exported_program: torch.export.ExportedProgram):
+    graph = exported_program.graph_module.graph
+    for node in graph.nodes:
+      if not (
+          node.op == "call_function"
+          and node.target == torch.ops.xla.mark_tensor.default
+      ):
+        continue
+
+      source, name, io_position, id, is_input = node.args[:5]
+      # Composite info:
+      # - name: odml.scaled_dot_product_attention
+      # - inputs: q, k, v, mask
+      if name == "odml.scaled_dot_product_attention" and is_input and io_position == 3:
+        if self.is_zero_tensor_node(source):
+          # Remove the mark_tensor call on the mask input by
+          # replacing the target with an identity function.
+          node.target = lambda *args, **kwargs: args[0]
+
+    exported_program.graph_module.graph.lint()
+    exported_program.graph_module.recompile()
+    return ExportedProgramPassResult(exported_program, True)

ai_edge_torch/generative/fx_passes/test/test_remove_sdpa_zero_mask_pass.py

@@ -0,0 +1,125 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import re
+from typing import Callable, Union
+import unittest
+
+import torch
+import torch_xla
+
+from ai_edge_torch.convert.fx_passes import CanonicalizePass
+from ai_edge_torch.convert.fx_passes import run_passes
+from ai_edge_torch.generative.fx_passes import RemoveSDPACompositeZeroMaskPass
+from ai_edge_torch.generative.layers.attention import SelfAttention
+import ai_edge_torch.generative.layers.model_config as layers_cfg
+import ai_edge_torch.generative.layers.unet.builder as unet_builder
+import ai_edge_torch.generative.layers.unet.model_config as unet_cfg
+
+
+def _export_to_stablehlo(func: Union[torch.nn.Module, Callable], export_args):
+  if not isinstance(func, torch.nn.Module):
+
+    class TestModule(torch.nn.Module):
+
+      def forward(self, *args, **kwargs):
+        return func(*args, **kwargs)
+
+    module = TestModule().eval()
+  else:
+    module = func
+
+  exported_program = torch.export.export(module, export_args)
+  exported_program = run_passes(
+      exported_program,
+      [
+          RemoveSDPACompositeZeroMaskPass(),
+          CanonicalizePass(),
+      ],
+  )
+
+  return torch_xla.stablehlo.exported_program_to_stablehlo(
+      exported_program
+  ).get_stablehlo_text()
+
+
+class TestRemoveSDPAZeroMaskPass(unittest.TestCase):
+
+  def test_self_attention_no_zero_mask_composite_input(self):
+    class SampleSdpaBlock(torch.nn.Module):
+      """Sample attention block with SDPA"""
+
+      def __init__(self, config: unet_cfg.AttentionBlock2DConfig):
+        super().__init__()
+        self.config = config
+        self.attention = SelfAttention(
+            config.attention_batch_size,
+            config.dim,
+            config.attention_config,
+            0,
+            enable_hlfb=config.enable_hlfb,
+        )
+
+      def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        B, C, H, W = input_tensor.shape
+        x = input_tensor
+        x = input_tensor.view(B, C, H * W)
+        x = x.transpose(-1, -2)
+        # x = x.contiguous()  # Prevent BATCH_MATMUL op in converted tflite.
+        x = self.attention(x)
+        x = x.transpose(-1, -2)
+        x = x.view(B, C, H, W)
+        return x
+
+    def get_model_config() -> unet_cfg.AttentionBlock2DConfig:
+      """Get configs for the Decoder of Stable Diffusion v1.5"""
+      in_channels = 3
+      latent_channels = 4
+      out_channels = 3
+      block_out_channels = [128, 256, 512, 512]
+      scaling_factor = 0.18215
+      layers_per_block = 3
+
+      norm_config = layers_cfg.NormalizationConfig(
+          layers_cfg.NormalizationType.GROUP_NORM, group_num=32
+      )
+
+      return unet_cfg.AttentionBlock2DConfig(
+          dim=block_out_channels[-1],
+          normalization_config=norm_config,
+          attention_config=layers_cfg.AttentionConfig(
+              num_heads=1,
+              num_query_groups=1,
+              qkv_use_bias=True,
+              output_proj_use_bias=True,
+              enable_kv_cache=False,
+              qkv_transpose_before_split=True,
+              rotary_percentage=0.0,
+          ),
+      )
+
+    stablehlo = _export_to_stablehlo(
+        SampleSdpaBlock(get_model_config()).eval(), (torch.rand(1, 512, 64, 64),)
+    )
+    print(stablehlo)
+    self.assertTrue(
+        re.search(
+            'stablehlo\.composite "odml\.scaled_dot_product_attention" %\d+, %\d+, %\d+ {',
+            stablehlo,
+        )
+    )
+
+
+if __name__ == '__main__':
+  unittest.main()