[cadence][g3] link m3 ops and add testcases (pytorch#8824)

Differential Revision: D71155209

Pull Request resolved: pytorch#9289

Author: zonglinpeng

backends/cadence/aot/functions_fusion_g3.yaml

@@ -51,7 +51,7 @@
 - op: clamp.Tensor_out
   kernels:
     - arg_meta: null
-      kernel_name: cadence::impl::G3::clamp_tensor_out
+      kernel_name: cadence::impl::G3::clamp_Tensor_out
 
 - op: clone.out
   kernels:
@@ -81,12 +81,12 @@
 - op: lt.Scalar_out
   kernels:
     - arg_meta: null
-      kernel_name: cadence::impl::G3::lt_scalar_out
+      kernel_name: cadence::impl::G3::lt_Scalar_out
 
 - op: lt.Tensor_out
   kernels:
     - arg_meta: null
-      kernel_name: cadence::impl::G3::lt_tensor_out
+      kernel_name: cadence::impl::G3::lt_Tensor_out
 
 - op: mul.out
   kernels:
@@ -155,7 +155,7 @@
 - op: where.self_out
   kernels:
     - arg_meta: null
-      kernel_name: cadence::impl::G3::where_out
+      kernel_name: cadence::impl::G3::where_self_out
 
 - op: native_layer_norm.out
   kernels:

backends/cadence/fusion_g3/operators/op_clamp.cpp

@@ -330,7 +330,7 @@ Tensor& clamp_out(
   return out;
 }
 
-Tensor& clamp_tensor_out(
+Tensor& clamp_Tensor_out(
     KernelRuntimeContext& ctx,
     const Tensor& in,
     const optional<Tensor>& min_opt,

backends/cadence/fusion_g3/operators/op_lt.cpp

@@ -24,7 +24,7 @@ namespace impl {
 namespace G3 {
 namespace native {
 
-Tensor& lt_tensor_out(
+Tensor& lt_Tensor_out(
     KernelRuntimeContext& ctx,
     const Tensor& a,
     const Tensor& b,
@@ -141,7 +141,7 @@ Tensor& lt_tensor_out(
   return out;
 }
 
-Tensor& lt_scalar_out(
+Tensor& lt_Scalar_out(
     KernelRuntimeContext& ctx,
     const Tensor& a,
     const Scalar& b,

backends/cadence/fusion_g3/operators/op_permute_copy.cpp

@@ -157,4 +157,4 @@ Tensor& permute_copy_out(
 } // namespace native
 } // namespace G3
 } // namespace impl
-} // namespace cadence
+} // namespace cadence

backends/cadence/fusion_g3/operators/op_where.cpp

@@ -24,7 +24,7 @@ namespace impl {
 namespace G3 {
 namespace native {
 
-Tensor& where_out(
+Tensor& where_self_out(
     KernelRuntimeContext& ctx,
     const Tensor& cond,
     const Tensor& a,

backends/cadence/fusion_g3/operators/targets.bzl

@@ -35,6 +35,14 @@ def define_operator(name: str, deps: list[str] | None = None) -> None:
 OPERATORS = [
     "add",
     "cat",
+    "clamp",
+    "lt",
+    "rsqrt",
+    "sigmoid",
+    "sqrt",
+    "tanh",
+    "transpose_copy",
+    "where",
     "dequantize",
     "mul",
     "native_layer_norm",

backends/cadence/utils/facto_util.py

@@ -22,9 +22,29 @@
 
 
 def apply_tensor_contraints(op_name: str, tensor_constraints: list[object]) -> None:
+    additional_tensor_constraints = [
+        cp.Dtype.In(lambda deps: [torch.int, torch.float]),
+        cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]),
+        cp.Value.Ge(lambda deps, dtype, struct: -(2**4)),
+        cp.Value.Le(lambda deps, dtype, struct: 2**4),
+        cp.Rank.Ge(lambda deps: 1),
+        cp.Size.Ge(lambda deps, r, d: 1),
+        cp.Size.Le(lambda deps, r, d: 2**9),
+    ]
+
     match op_name:
+        case "where.self":
+            additional_tensor_constraints = [
+                cp.Dtype.In(lambda deps: [torch.float, torch.int, torch.bool]),
+                cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]),
+                cp.Value.Ge(lambda deps, dtype, struct: -(2**4)),
+                cp.Value.Le(lambda deps, dtype, struct: 2**4),
+                cp.Rank.Ge(lambda deps: 1),
+                cp.Size.Ge(lambda deps, r, d: 1),
+                cp.Size.Le(lambda deps, r, d: 2**9),
+            ]
         case "sigmoid.default" | "rsqrt.default":
-            tensor_constraints.extend(
+            additional_tensor_constraints.extend(
                 [
                     cp.Dtype.In(lambda deps: [torch.float]),
                     cp.Rank.Le(lambda deps: 2**2),
@@ -33,45 +53,35 @@ def apply_tensor_contraints(op_name: str, tensor_constraints: list[object]) -> N
                 ]
             )
         case "mean.dim":
-            tensor_constraints.extend(
+            additional_tensor_constraints.extend(
                 [
                     cp.Dtype.In(lambda deps: [torch.float]),
                     cp.Rank.Le(lambda deps: 2**2),
                 ]
             )
         case "exp.default":
-            tensor_constraints.extend(
+            additional_tensor_constraints.extend(
                 [
                     cp.Rank.Le(lambda deps: 2**3),
                     cp.Value.Ge(lambda deps, dtype, struct: -(2**2)),
                     cp.Value.Le(lambda deps, dtype, struct: 2**2),
                 ]
             )
         case "slice_copy.Tensor":
-            tensor_constraints.extend(
+            additional_tensor_constraints.extend(
                 [
                     cp.Rank.Le(lambda deps: 2),
                     cp.Value.Ge(lambda deps, dtype, struct: 1),
                     cp.Value.Le(lambda deps, dtype, struct: 2),
                 ]
             )
         case _:
-            tensor_constraints.extend(
+            additional_tensor_constraints.extend(
                 [
                     cp.Rank.Le(lambda deps: 2**2),
                 ]
             )
-    tensor_constraints.extend(
-        [
-            cp.Dtype.In(lambda deps: [torch.int, torch.float]),
-            cp.Dtype.NotIn(lambda deps: [torch.int64, torch.float64]),
-            cp.Value.Ge(lambda deps, dtype, struct: -(2**4)),
-            cp.Value.Le(lambda deps, dtype, struct: 2**4),
-            cp.Rank.Ge(lambda deps: 1),
-            cp.Size.Ge(lambda deps, r, d: 1),
-            cp.Size.Le(lambda deps, r, d: 2**9),
-        ]
-    )
+    tensor_constraints.extend(additional_tensor_constraints)
 
 
 def apply_scalar_contraints(op_name: str) -> list[ScalarDtype]: