experimental initial torch-tfl setup

PiperOrigin-RevId: 736573537

Author: chunnienc

ai_edge_torch/odml_torch/experimental/torch_tfl/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch-TFL ops definitions, decompositions, and lowerings."""
+from ai_edge_torch.odml_torch.experimental.torch_tfl import _decomps
+from ai_edge_torch.odml_torch.experimental.torch_tfl import _lowerings
+from ai_edge_torch.odml_torch.experimental.torch_tfl import _ops
+
+decomps = _decomps.decomps

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -0,0 +1,38 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch ops to Torch-TFL decompositions."""
+from ai_edge_torch.odml_torch.experimental.torch_tfl import _ops
+import torch
+
+decomps = {}
+
+
+def register_decomp(op):
+  global decomps
+  ops = [op]
+  if isinstance(op, torch._ops.OpOverloadPacket):
+    ops = [getattr(op, overload) for overload in op.overloads()]
+
+  def register(decomp_fn):
+    for op in ops:
+      decomps[op] = decomp_fn
+    return decomp_fn
+
+  return register
+
+
+@register_decomp(torch.ops.aten.mm.default)
+def _aten_mm_decomp(x, y):
+  return torch.ops.tfl.batch_matmul(x, y)

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -0,0 +1,18 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch-TFL op to MLIR lowerings."""
+from ai_edge_torch.odml_torch.lowerings import registry
+
+lower = registry.lower

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -0,0 +1,40 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch-TFL op definitions and fake implementations."""
+from ai_edge_torch.odml_torch.experimental.torch_tfl import torch_library_utils
+import torch
+
+custom_op_with_fake = torch_library_utils.custom_op_with_fake
+
+
+@custom_op_with_fake("tfl::batch_matmul")
+def tfl_batch_matmul(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    adj_x: bool = False,
+    adj_y: bool = False,
+    asymmetric_quantize_inputs: bool = False,
+) -> torch.Tensor:
+  if asymmetric_quantize_inputs:
+    raise NotImplementedError(
+        "asymmetric_quantize_inputs=True is not implemented"
+    )
+  if x.ndim < 2 or y.ndim < 2:
+    raise ValueError("Input tensors must have at least 2 dimensions.")
+  if adj_x:
+    x = torch.transpose(x, -1, -2)
+  if adj_y:
+    y = torch.transpose(y, -1, -2)
+  return torch.matmul(x, y)

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -0,0 +1,91 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Numerical validation tests for torch ops and Torch-TFL ops."""
+from ai_edge_torch import testing
+from ai_edge_torch.odml_torch.experimental import torch_tfl
+import torch
+from torch.utils import _pytree as pytree
+
+from absl.testing import absltest as googletest
+from absl.testing import parameterized
+
+export_with_tensor_inputs_only = testing.export_with_tensor_inputs_only
+
+
+def rnd(dtype, shape, min_v=None, max_v=None):
+  """Shortcut for creating a random torch tensor."""
+  if dtype in (torch.int32, torch.int64, torch.bool):
+    min_v = min_v if min_v else 1
+    max_v = max_v if max_v else 10
+    return torch.randint(min_v, max_v, shape).to(dtype)
+  else:
+    min_v = min_v if min_v else 0.0
+    max_v = max_v if max_v else 1.0
+    return (torch.rand(shape) * (max_v - min_v) + min_v).to(dtype)
+
+
+class TestTorchTFLImpls(parameterized.TestCase):
+  """Numerical validation tests for torch ops and Torch-TFL ops.
+
+  The op test suite is forked from
+  ai_edge_torch/odml_torch/test/test_core_aten_ops.py. Eventually, we should
+  merge the two test suites.
+  """
+
+  def setUp(self):
+    super().setUp()
+    torch.manual_seed(0)
+
+  def _assert_export_and_close(
+      self, func, args, kwargs, atol=1e-3, rtol=1e-5, equal_nan=True
+  ):
+    """Assert func, args, and kwargs can be lowered and pass numerical validation."""
+    with self.subTest("torch_eval"):
+      expected = func(*args, **kwargs)
+
+      with self.subTest("export_and_decompse"):
+        exported_program = export_with_tensor_inputs_only(func, args, kwargs)
+        exported_program = exported_program.run_decompositions(
+            torch_tfl.decomps
+        )
+
+        with self.subTest("decomp_eval"):
+          args, kwargs = exported_program.example_inputs
+          actual = exported_program.module()(*args, **kwargs)
+
+          with self.subTest("torch_lower_eval_diff:" + str(atol)):
+            expected_flat, expected_spec = pytree.tree_flatten(expected)
+            actual_flat, actual_spec = pytree.tree_flatten(actual)
+
+            self.assertEqual(expected_spec, actual_spec)
+            for v1, v2 in zip(expected_flat, actual_flat):
+              torch.testing.assert_close(
+                  v1, v2, atol=atol, rtol=rtol, equal_nan=equal_nan
+              )
+
+  @parameterized.named_parameters(
+      # fmt: off
+      # pyformat: disabledef
+      ("aten_mm_0", torch.ops.aten.mm.default, (rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
+      ("aten_mm_1", torch.ops.aten.mm.default, (rnd(torch.float32, (2, 10)), rnd(torch.float32, (10, 5)),), dict()),
+      # fmt: on
+      # pyformat: enable
+  )
+  def test_op(self, op, args, kwargs):
+    self._assert_export_and_close(op, args, kwargs)
+
+
+if __name__ == "__main__":
+  googletest.main()

ai_edge_torch/odml_torch/experimental/torch_tfl/torch_library_utils.py

@@ -0,0 +1,37 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Utility functions for defining custom ops in torch library."""
+from typing import Callable, Iterable
+import torch
+
+
+def custom_op_with_fake(
+    name: str,
+    *,
+    mutates_args: str | Iterable[str] = (),
+    schema: str | None = None,
+):
+  """Defines a custom op with a FakeTensor implementation using the same function."""
+
+  def register(fn: Callable[..., object]):
+    op = torch.library.custom_op(
+        name,
+        mutates_args=mutates_args,
+        schema=schema,
+    )(fn)
+    torch.library.register_fake(name)(fn)
+    return op
+
+  return register

ai_edge_torch/testing/__init__.py

@@ -12,3 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from ai_edge_torch.testing import export
+from ai_edge_torch.testing import model_coverage
+
+export_with_tensor_inputs_only = export.export_with_tensor_inputs_only

ai_edge_torch/testing/export.py

@@ -0,0 +1,83 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch export utilities for testing."""
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch.utils import _pytree as pytree
+
+
+def export_with_tensor_inputs_only(
+    model: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> torch.export.ExportedProgram:
+  """Exports a PyTorch model, treating only tensor inputs as export inputs.
+
+  This function takes a PyTorch model and its input arguments (positional and
+  keyword) and exports it using `torch.export.export`. However, it modifies
+  the export process such that only the `torch.Tensor` arguments in the
+  inputs are considered as export inputs to the exported graph. All other
+  argument types (e.g., scalars, lists, tuples containing non-tensors) are
+  treated as constants.
+
+  This is useful for testing scenarios where you want to export a model but
+  want to avoid issues that might arise from non-tensor inputs
+  being treated as variables, or when you specifically want to focus on the
+  graph structure based on tensor operations.
+
+  Args:
+    model: The PyTorch `nn.Module` to be exported.
+    args: A tuple of positional arguments to be passed to the model's `forward`
+      method.
+    kwargs: A dictionary of keyword arguments to be passed to the model's
+      `forward` method.
+
+  Returns:
+    torch.export.ExportedProgram: The exported program representing the model
+    computation with only tensor inputs being export inputs.
+  """
+  flatten_args, treespec = pytree.tree_flatten([args, kwargs])
+
+  export_args = []
+  indices = []
+  for i, arg in enumerate(flatten_args):
+    if isinstance(arg, torch.Tensor):
+      export_args.append(arg)
+      indices.append(i)
+
+  class ModuleWrapper(torch.nn.Module):
+
+    def __init__(self, func, original_args, original_kwargs):
+      super().__init__()
+      self.original_args = list(flatten_args)
+      self.func = func
+
+    def forward(self, *export_args):
+      flatten_args = self.original_args.copy()
+      for i, arg in zip(indices, export_args):
+        flatten_args[i] = arg
+      args, kwargs = pytree.tree_unflatten(flatten_args, treespec)
+      return self.func(*args, **kwargs)
+
+  export_args = tuple(export_args)
+  export_kwargs = {}
+  return torch.export.export(
+      ModuleWrapper(model, args, kwargs).eval(),
+      export_args,
+      export_kwargs,
+  )