Make custom op syntax dependent on torch version

Signed-off-by: Andrea Fasoli <[email protected]>

Author: andrea-fasoli

fms_mo/aiu_addons/i8i8/i8i8_aiu_op.py

@@ -17,18 +17,42 @@
 import logging
 
 # Third Party
+from packaging.version import Version
 import torch
 import torch.nn.functional as F
 
-logger = logging.getLogger(__name__)
-
 # pylint: disable=unused-argument
 # i8i8 op must be registered with specific I/O, even if not in use by the op function
 
 # pylint: disable=not-callable
 # torch.nn.functional.linear not recognized as callable
 # open issue in PyLint: https://github.com/pytorch/pytorch/issues/119482
 
+logger = logging.getLogger(__name__)
+torch_version = Version(torch.__version__.split("+", maxsplit=1)[0])
+
+
+def implement_op_decorator(pt_ver, op_namespace_id):
+    """Version-dependent decorator for custom op implementation."""
+
+    def decorator(func):
+        if pt_ver <= Version("2.4"):
+            return torch.library.impl(op_namespace_id, "default")(func)
+        return torch.library.custom_op(op_namespace_id, mutates_args=())(func)
+
+    return decorator
+
+
+def register_op_decorator(pt_ver, op_namespace_id):
+    """Version-dependent decorator for custom op registration."""
+
+    def decorator(func):
+        if pt_ver <= Version("2.4"):
+            return torch.library.impl_abstract(op_namespace_id)(func)
+        return torch.library.register_fake(op_namespace_id)(func)
+
+    return decorator
+
 
 def register_aiu_i8i8_op():
     """Register AIU-specific op to enable torch compile without graph break.
@@ -42,8 +66,16 @@ def register_aiu_i8i8_op():
         logger.warning("AIU op has already been registered")
         return
     op_namespace_id = "fms_mo::i8i8_aiu"
+    if torch_version <= Version("2.4"):
+        torch.library.define(
+            op_namespace_id,
+            "(Tensor x, Tensor weight, Tensor bias, Tensor qdata, "
+            "str weight_quant_type, str activ_quant_type, "
+            "bool smoothquant) "
+            "-> Tensor",
+        )
 
-    @torch.library.custom_op(op_namespace_id, mutates_args=())
+    @implement_op_decorator(torch_version, op_namespace_id)
     def i8i8_aiu(
         x: torch.Tensor,
         weight: torch.Tensor,
@@ -78,7 +110,7 @@ def i8i8_aiu(
 
         return F.linear(x_dq.to(dtype), w_dq.to(dtype), bias.to(dtype))
 
-    @torch.library.register_fake(op_namespace_id)
+    @register_op_decorator(torch_version, op_namespace_id)
     def _(x, weight, bias, qdata, weight_quant_type, activ_quant_type, smoothquant):
         """OP template of I/O sizes"""