[Relax][PyTorch] Add PReLU Op Support for Exported Program and FX graph (apache#17816)

* prelu op support and test script added

* end-of-file issue fixed

* trailing whitespace issue fixed

* fixing lint issues

* fix assertion error in test_op_nn.py file

* add test script in test_frontend_nn_op.py

* include wrapper function for prelu in op.py

* fixing unity check issue by modifying test func

* conflicts resolved

* add doc for prelu op axis arg

* fixed failing checks issue

---------

Co-authored-by: deivanayakisankaralingam <deiva@Deivanayaki>

Author: 2 people

include/tvm/relax/attrs/nn.h

@@ -468,6 +468,15 @@ struct SoftplusAttrs : public tvm::AttrsNode<SoftplusAttrs> {
   }
 };
 
+/*! \brief Attributes used in PReLU operator */
+struct PReluAttrs : public tvm::AttrsNode<PReluAttrs> {
+  int axis;
+
+  TVM_DECLARE_ATTRS(PReluAttrs, "relax.attrs.PReluAttrs") {
+    TVM_ATTR_FIELD(axis).describe("The axis along which the alpha values are applied.");
+  }
+};
+
 /*! \brief Attributes used in batch_norm operator */
 struct BatchNormAttrs : public tvm::AttrsNode<BatchNormAttrs> {
   int axis;

python/tvm/relax/frontend/nn/op.py

@@ -1072,6 +1072,34 @@ def softplus(x: Tensor, beta: float = 1.0, threshold: float = 20.0, name: str =
     return wrap_nested(_op.nn.softplus(x._expr, beta=beta, threshold=threshold), name)
 
 
+def prelu(x: Tensor, alpha: Tensor, name: str = "prelu"):
+    r"""Parametric ReLU activation function.
+
+    .. math::
+        \text{PReLU}(x) = \begin{cases}
+            x & \text{if } x \geq 0 \\
+            \alpha \cdot x & \text{if } x < 0
+        \end{cases}
+
+    Parameters
+    ----------
+    x : Tensor
+        The input data.
+
+    alpha : Tensor
+        Slope coefficient for the negative part of the input.
+
+    name : str, optional
+        Optional name for the operation. Default is "prelu".
+
+    Returns
+    -------
+    result : Tensor
+        The computed result.
+    """
+    return wrap_nested(_op.nn.prelu(x._expr, alpha._expr), name)
+
+
 def tanh(x: Tensor, name: str = "tanh") -> Tensor:
     r"""Applies the hyperbolic tangent function.
 

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -307,6 +307,12 @@ def _log_softmax(self, node: fx.Node) -> relax.Var:
         dim = node.args[1] if len(node.args) > 1 else node.kwargs.get("dim", -1)
         return self.block_builder.emit(relax.op.nn.log_softmax(x, dim))
 
+    def _prelu(self, node: fx.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        alpha = self.env[node.args[1]]
+        axis = 0 if len(x.struct_info.shape) == 1 else 1
+        return self.block_builder.emit(relax.op.nn.prelu(x, alpha, axis))
+
     def _round(self, node: fx.Node) -> relax.Expr:
         if node.kwargs.get("decimals", 0) != 0:
             raise ValueError("specifying decimals for round is not supported yet")

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -299,6 +299,7 @@ def create_convert_map(
             "log1p.default": self._log1p,
             "log_softmax.int": self._log_softmax,
             "neg.default": self._unary_op(relax.op.negative),
+            "prelu.default": self._prelu,
             "reciprocal.default": self._reciprocal,
             "relu.default": self._unary_op(relax.op.nn.relu),
             "relu_.default": self._unary_op(relax.op.nn.relu),

python/tvm/relax/frontend/torch/fx_translator.py

@@ -103,6 +103,14 @@ def _log_softmax_module(self, node: fx.Node) -> relax.Var:
         assert dim is not None
         return self.block_builder.emit(relax.op.nn.log_softmax(x, dim))
 
+    def _prelu_module(self, node: fx.Node) -> relax.Var:
+        x = self.env[node.args[0]]
+        module = self.named_modules[node.target]
+        alpha_tensor = module.weight.numpy()
+        alpha = relax.const(alpha_tensor, dtype="float32")
+        axis = 0 if len(x.struct_info.shape) == 1 else 1  # Extract Channel size
+        return self.block_builder.emit(relax.op.nn.prelu(x, alpha, axis))
+
     def _softmax_module(self, node: fx.Node) -> relax.Var:
         x = self.env[node.args[0]]
         module = self.named_modules[node.target]
@@ -595,6 +603,7 @@ def create_convert_map(
             nn.Identity: lambda node: self.env[node.args[0]],
             nn.LeakyReLU: self._leakyrelu_module,
             nn.LogSoftmax: self._log_softmax_module,
+            nn.PReLU: self._prelu_module,
             nn.ReLU: self._unary_op(relax.op.nn.relu),
             nn.ReLU6: lambda node: self.block_builder.emit(
                 relax.op.clip(self.env[node.args[0]], 0, 6)
@@ -657,6 +666,7 @@ def create_convert_map(
             "logical_not": self._unary_op(relax.op.logical_not),
             "log_softmax": self._log_softmax,
             "neg": self._unary_op(relax.op.negative),
+            "prelu": self._prelu,
             "reciprocal": self._reciprocal,
             "relu": self._unary_op(relax.op.nn.relu),
             "round": self._round,

python/tvm/relax/op/nn/__init__.py

@@ -43,6 +43,7 @@
     max_pool3d,
     nll_loss,
     pad,
+    prelu,
     relu,
     rms_norm,
     selu,

python/tvm/relax/op/nn/nn.py

@@ -1431,6 +1431,32 @@ def log_softmax(data: Expr, axis: int = -1) -> Expr:
     return _ffi_api.log_softmax(data, axis)  # type: ignore
 
 
+def prelu(data: Expr, alpha: Expr, axis: int = 1) -> Expr:
+    r"""Parametric Rectified Linear Unit (PReLU).
+
+    .. math::
+        PReLU(x) = x \text{ if } x > 0 \text{ else } \alpha * x
+
+    Parameters
+    ----------
+    data : relax.Expr
+        The input tensor.
+
+    alpha : relax.Expr
+        The learnable slope tensor, applied channel-wise.
+
+    axis : int
+        The axis along which the `alpha` values are applied
+        Default is 1 (assuming NCHW format).
+
+    Returns
+    -------
+    result : relax.Expr
+        The computed result.
+    """
+    return _ffi_api.prelu(data, alpha, axis)
+
+
 def batch_norm(
     data: Expr,
     gamma: Expr,

python/tvm/relax/transform/legalize_ops/nn.py

@@ -469,6 +469,11 @@ def _nn_leakyrelu(bb: BlockBuilder, call: Call) -> Expr:
     return bb.call_te(topi.nn.leaky_relu, call.args[0], call.attrs.alpha)
 
 
+@register_legalize("relax.nn.prelu")
+def _nn_prelu(bb: BlockBuilder, call: Call) -> Expr:
+    return bb.call_te(topi.nn.prelu, call.args[0], call.args[1], call.attrs.axis)
+
+
 @register_legalize("relax.nn.gelu")
 def _nn_gelu(bb: BlockBuilder, call: Call) -> Expr:
     def te_gelu(x: te.Tensor):

python/tvm/topi/nn/elemwise.py

@@ -129,6 +129,9 @@ def prelu(x, slope, axis=1):
 
     assert len(slope.shape) == 1
     assert axis < len(x.shape)
+    slope = te.compute(
+        (get_const_int(x.shape[axis]),), lambda c: slope[0], name="slope_broadcasted"
+    )
     assert get_const_int(slope.shape[0]) == get_const_int(x.shape[axis])
 
     def _compute_channelwise(*indices):

src/relax/op/nn/nn.cc

@@ -81,6 +81,27 @@ TVM_REGISTER_OP("relax.nn.softplus")
                                 InferStructInfoUnaryArith</*require_float_dtype=*/true>)
     .set_attr<Bool>("FPurity", Bool(true));
 
+/* relax.nn.prelu */
+TVM_REGISTER_NODE_TYPE(PReluAttrs);
+
+Expr prelu(Expr data, Expr alpha, int axis = 1) {
+  auto attrs = make_object<PReluAttrs>();
+  attrs->axis = axis;
+  static const Op& op = Op::Get("relax.nn.prelu");
+  return Call(op, {data, alpha}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relax.op.nn.prelu").set_body_typed(prelu);
+
+TVM_REGISTER_OP("relax.nn.prelu")
+    .set_num_inputs(2)
+    .add_argument("data", "Tensor", "The input tensor.")
+    .add_argument("alpha", "Tensor", "The channel-wise learnable slope.")
+    .set_attrs_type<PReluAttrs>()
+    .set_attr<FInferStructInfo>("FInferStructInfo",
+                                InferStructInfoUnaryArith</*require_float_dtype=*/true>)
+    .set_attr<Bool>("FPurity", Bool(true));
+
 /* relax.nn.softmax */
 TVM_REGISTER_NODE_TYPE(SoftmaxAttrs);