move llama.fallback out from model_sharding

Author: chenweng-quic

exir/passes/spec_prop_pass.py

@@ -9,7 +9,8 @@
 import operator
 from typing import Optional
 
-import executorch.extension.llm.custom_ops.model_sharding  # noqa: F401
+# register llama.fallback
+import executorch.extension.llm.custom_ops.op_fallback  # noqa: F401
 
 import torch
 from executorch.exir.delegate import executorch_call_delegate

extension/llm/custom_ops/model_sharding.py

@@ -7,8 +7,9 @@
 import re
 from typing import List
 
-import torch
+import executorch.extension.llm.custom_ops.op_fallback  # noqa: F401
 
+import torch
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_PASS_ACTIVATE_KEY,
     QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY,
@@ -17,27 +18,6 @@
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.export.exported_program import ExportedProgram
-from torch.library import impl, Library
-
-
-fallback_op_lib = Library("llama", "DEF")
-# registering an operator.
-fallback_op_lib.define("fallback(Tensor input) -> Tensor")
-
-
-@impl(fallback_op_lib, "fallback")
-def fallback_impl(a: torch.Tensor) -> torch.Tensor:
-    return a
-
-
-# registering the out variant.
-fallback_op_lib.define("fallback.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)")
-
-
-@impl(fallback_op_lib, "fallback.out")
-def fallback_out_impl(a: torch.Tensor, *, out: torch.Tensor) -> torch.Tensor:
-    out.copy_(a)
-    return out
 
 
 class SplitGraph(ExportPass):

extension/llm/custom_ops/op_fallback.py

@@ -0,0 +1,29 @@
+# Copyright (c) Qualcomm Innovation Center, Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# pyre-ignore-all-errors
+
+import torch
+
+from torch.library import impl, Library
+
+fallback_op_lib = Library("llama", "DEF")
+# registering an operator.
+fallback_op_lib.define("fallback(Tensor input) -> Tensor")
+
+
+@impl(fallback_op_lib, "fallback")
+def fallback_impl(a: torch.Tensor) -> torch.Tensor:
+    return a
+
+
+# registering the out variant.
+fallback_op_lib.define("fallback.out(Tensor input, *, Tensor(a!) output) -> Tensor(a!)")
+
+
+@impl(fallback_op_lib, "fallback.out")
+def fallback_out_impl(a: torch.Tensor, *, out: torch.Tensor) -> torch.Tensor:
+    out.copy_(a)
+    return out