[Op-Authoring] Adding mapping from torch ops to ExprHandles (#205)

* [Op-Authoring] Adding mapping from torch ops to ExprHandles

Dependent on PR - pytorch/pytorch#66612

* Cleaing up and bug fix

* Retrigger cI

* Retirgger CI

* Retrigger CI

Author: anijain2305

functorch/_src/operator_authoring.py

@@ -3,6 +3,7 @@
 import inspect
 import itertools
 from typing import Callable, List, Union, Tuple, Optional
+import operator
 
 import torch
 from torch import fx
@@ -12,6 +13,51 @@
 FOLD_ALIASES = True
 _SHAPE_TYPES = {"one", "other"}
 _STRIDE_TYPES = {"zero", "one", "contiguous", "transposed_contiguous", "as_arg"}
+_identity = lambda x: x
+_TORCH_TO_EXPR_MAP = {
+    "sin": _te.sin,
+    "cos": _te.cos,
+    "tan": _te.tan,
+    "asin": _te.asin,
+    "acos": _te.acos,
+    "atan": _te.atan,
+    "sinh": _te.sinh,
+    "cosh": _te.cosh,
+    "tanh": _te.tanh,
+    "sigmoid": _te.sigmoid,
+    "exp": _te.exp,
+    "expm1": _te.expm1,
+    "abs": _te.abs,
+    "log": _te.log,
+    "log2": _te.log2,
+    "log10": _te.log10,
+    "log1p": _te.log1p,
+    "erf": _te.erf,
+    "erfc": _te.erfc,
+    "sqrt": _te.sqrt,
+    "rsqrt": _te.rsqrt,
+    "ceil": _te.ceil,
+    "floor": _te.floor,
+    "round": _te.round,
+    "trunc": _te.trunc,
+    "frac": _te.frac,
+    "lgamma": _te.lgamma,
+    "isnan": _te.isnan,
+    "add": operator.add,
+    "sub": operator.sub,
+    "subtract": operator.sub,
+    "mul": operator.mul,
+    "multiply": operator.mul,
+    "divide": operator.truediv,
+    "div": operator.truediv,
+    "remainder": _te.remainder,
+    "fmod": _te.fmod,
+    "pow": _te.pow,
+    "atan2": _te.atan2,
+    "detach": _identity,
+    "neg": lambda x: _create_constant(0.0, torch.float32) - x,
+}
+
 _int = _te.ExprHandle.int
 
 
@@ -40,8 +86,8 @@ def _combine_dtype(a: torch.dtype, b: torch.dtype):
     ).dtype
 
 
-def _fx_replace_constants(fn: Callable, dtype: torch.dtype):
-    """Convert the constants in the user function to TensorExpr constants"""
+def _fx_to_expr(fn: Callable, dtype: torch.dtype):
+    """Convert the fx graph to equivalent Tensor Expr"""
 
     def apply(arg):
         if isinstance(arg, (int, float)):
@@ -52,6 +98,27 @@ def apply(arg):
     for node in list(gm.graph.nodes):
         with gm.graph.inserting_before(node):
             node.args = tuple(apply(a) for a in node.args)
+            if node.op == "call_function":
+                if node.target.__name__ not in _TORCH_TO_EXPR_MAP:
+                    raise NotImplementedError(
+                        "Missing mapping from op ",
+                        node.target.__name__,
+                        " to Tensor Expr",
+                    )
+
+                    # Get the parser function to parse the torch op to tensor expr handle
+
+                def _parser(*args, op_name):
+                    return _TORCH_TO_EXPR_MAP[op_name](*args)
+
+                new_node = gm.graph.create_node(
+                    "call_function",
+                    _parser,
+                    node.args,
+                    {"op_name": node.target.__name__},
+                )
+                node.replace_all_uses_with(new_node)
+                gm.graph.erase_node(node)
     gm.recompile()
     return gm
 
@@ -298,7 +365,7 @@ def compute_code(self):
             _te.Cast.make(self.dtype, buf.load(self.indexing(stride)))
             for buf, stride in zip(input_bufs, input_strides)
         ]
-        val = _fx_replace_constants(self.pointwise_fn, self.dtype)(*inputs)
+        val = _fx_to_expr(self.pointwise_fn, self.dtype)(*inputs)
         out = _te.Block(
             [
                 buf.store(self.indexing(stride), val)

test/test_operator_authoring.py

@@ -126,6 +126,119 @@ def example(x):
         x = torch.randn(8, device=self.device)
         torch.testing.assert_allclose(x + 3, graph(x))
 
+    def test_unary_ops(self):
+        unary_operators = [
+            torch.sin,
+            torch.cos,
+            torch.tan,
+            torch.asin,
+            torch.acos,
+            torch.atan,
+            torch.sinh,
+            torch.cosh,
+            torch.tanh,
+            torch.sigmoid,
+            torch.exp,
+            torch.expm1,
+            torch.abs,
+            torch.log,
+            torch.log2,
+            torch.log10,
+            torch.log1p,
+            torch.erf,
+            torch.erfc,
+            torch.sqrt,
+            torch.rsqrt,
+            torch.ceil,
+            torch.floor,
+            torch.round,
+            torch.trunc,
+            torch.lgamma,
+            torch.ops.aten.sin,
+            torch.ops.aten.cos,
+            torch.ops.aten.tan,
+            torch.ops.aten.asin,
+            torch.ops.aten.acos,
+            torch.ops.aten.atan,
+            torch.ops.aten.sinh,
+            torch.ops.aten.cosh,
+            torch.ops.aten.tanh,
+            torch.ops.aten.sigmoid,
+            torch.ops.aten.exp,
+            torch.ops.aten.expm1,
+            torch.ops.aten.abs,
+            torch.ops.aten.log,
+            torch.ops.aten.log2,
+            torch.ops.aten.log10,
+            torch.ops.aten.log1p,
+            torch.ops.aten.erf,
+            torch.ops.aten.erfc,
+            torch.ops.aten.sqrt,
+            torch.ops.aten.rsqrt,
+            torch.ops.aten.ceil,
+            torch.ops.aten.floor,
+            torch.ops.aten.round,
+            torch.ops.aten.trunc,
+            torch.ops.aten.lgamma,
+            # TODO - Failure in generating loop nests here for the following ops
+            # torch.frac,
+            # torch.isnan,
+        ]
+
+        for unary_op in unary_operators:
+            fn = lambda x: unary_op(x)
+            pointwise_fn = pointwise_operator(fn)
+            a = torch.rand(2, 3)
+            ref = fn(a)
+            res = pointwise_fn(a)
+            assert torch.allclose(ref, res, atol=1e-3, rtol=1e-3)
+
+    def test_binary_ops(self):
+        binary_operators = [
+            torch.add,
+            torch.sub,
+            torch.subtract,
+            torch.mul,
+            torch.multiply,
+            torch.divide,
+            torch.div,
+            torch.fmod,
+            torch.pow,
+            torch.atan2,
+            # torch.remainder, #TODO - Fails allclose check
+            torch.ops.aten.add,
+            torch.ops.aten.sub,
+            torch.ops.aten.subtract,
+            torch.ops.aten.mul,
+            torch.ops.aten.multiply,
+            torch.ops.aten.divide,
+            torch.ops.aten.div,
+            torch.ops.aten.fmod,
+            torch.ops.aten.pow,
+            torch.ops.aten.atan2,
+        ]
+        for binary_op in binary_operators:
+            fn = lambda x, y: binary_op(x, y)
+            pointwise_fn = pointwise_operator(fn)
+            a = torch.rand(2, 3)
+            b = torch.rand(2, 3)
+            ref = fn(a, b)
+            res = pointwise_fn(a, b)
+            assert torch.allclose(ref, res, atol=1e-3, rtol=1e-3)
+
+    def test_bias_gelu(self):
+        def bias_gelu(bias, y):
+            x = bias + y
+            return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+        bias = torch.rand(1, 768)
+        y = torch.rand(64, 768)
+        ref = bias_gelu(bias, y)
+
+        pointwise_fn = pointwise_operator(bias_gelu)
+        res = pointwise_fn(bias, y)
+        assert torch.allclose(ref, res, atol=1e-3, rtol=1e-3)
+
 
 @unittest.skipIf(not HAS_CUDA, "GPU tests require CUDA")
 class TestOperatorAuthoringGPU(TestOperatorAuthoring):