enable tfl direct lowering on shlo custom_call carrier

PiperOrigin-RevId: 736594309

Author: chunnienc

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -13,6 +13,52 @@
 # limitations under the License.
 # ==============================================================================
 """Torch-TFL op to MLIR lowerings."""
+from typing import Sequence
+from ai_edge_torch import odml_torch
+from ai_edge_torch.odml_torch.experimental.torch_tfl import _ops
 from ai_edge_torch.odml_torch.lowerings import registry
+from ai_edge_torch.odml_torch.lowerings import utils as lowering_utils
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import torch
 
 lower = registry.lower
+LoweringContext = odml_torch.lowerings.context.LoweringContext
+
+
+def _ir_operation(
+    name: str,
+    results: Sequence[ir.Type],
+    operands: Sequence[ir.Value] | None = None,
+    attributes: dict[str, ir.Attribute] | None = None,
+):
+  """Helper function to create an IR operation in StableHLO CustomCall carrier."""
+  attributes = ir.DictAttr.get(attributes)
+  return stablehlo.custom_call(
+      result=results,
+      inputs=operands,
+      call_target_name=ir.StringAttr.get(name),
+      has_side_effect=ir.BoolAttr.get(False),
+      backend_config=attributes,
+      api_version=ir.IntegerAttr.get(ir.IntegerType.get_signless(32), 4),
+  )
+
+
+@lower(torch.ops.tfl.batch_matmul.default)
+def _tfl_batch_matmul_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    y: ir.Value,
+    adj_x: bool = False,
+    adj_y: bool = False,
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.batch_matmul",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x, y],
+      attributes={
+          "adj_x": ir.BoolAttr.get(adj_x),
+          "adj_y": ir.BoolAttr.get(adj_y),
+          "asymmetric_quantize_inputs": ir.BoolAttr.get(False),
+      },
+  )

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -21,16 +21,8 @@
 
 @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,
+    x: torch.Tensor, y: torch.Tensor, adj_x: bool = False, adj_y: 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:

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -18,14 +18,16 @@
 import numbers
 from typing import Any
 from typing import Optional
-
+from ai_edge_torch.odml_torch import export_utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
 import numpy as np
 import torch
+import torch.utils._pytree as pytree
 
 
-def torch_dtype_to_ir_element_type(dtype):
+def torch_dtype_to_ir_element_type(dtype) -> ir.Type:
+  """Builds ir.Type from torch dtype."""
   ty_get = {
       torch.double: ir.F64Type.get,
       torch.float32: ir.F32Type.get,
@@ -39,6 +41,27 @@ def torch_dtype_to_ir_element_type(dtype):
   return ty_get()
 
 
+def node_meta_to_ir_types(node: torch.fx.Node) -> list[ir.Type]:
+  """Builds IR result types from torch FX node meta."""
+  tensor_meta = node.meta.get("tensor_meta") or node.meta.get("val")
+  if not tensor_meta:
+    raise RuntimeError(f"{node.name} does not have tensor meta")
+
+  tensor_meta_list, _ = pytree.tree_flatten(
+      [tensor_meta],
+      is_leaf=lambda x: hasattr(x, "dtype") and hasattr(x, "shape"),
+  )
+  results = []
+  for meta in tensor_meta_list:
+    shape = [
+        export_utils.IR_DYNAMIC if export_utils.is_torch_dynamic(dim) else dim
+        for dim in meta.shape
+    ]
+    elty = torch_dtype_to_ir_element_type(meta.dtype)
+    results.append(ir.RankedTensorType.get(shape, elty))
+  return results
+
+
 def splat(val, ty, shape=tuple(), *, loc: Optional[Any] = None):
   if isinstance(ty, ir.IntegerType):
     if ty.width == 1: