Finish T1-1-18: group_norm、reshape、bitwise_xor、isclose、rrelu

Author: PPPoint-t

include/infinicore/ops/bitwise_xor.hpp

@@ -0,0 +1,17 @@
+#pragma once
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class BitwiseXor {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor a, Tensor b, Tensor out);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor bitwise_xor(Tensor a, Tensor b);
+Tensor& bitwise_xor_out(Tensor a, Tensor b, Tensor& out);
+
+} // namespace infinicore::op

include/infinicore/ops/group_norm.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+class GroupNorm {
+public:
+    using schema = void (*)(Tensor, int64_t, std::optional<Tensor>, std::optional<Tensor>, double, Tensor);
+    static void execute(Tensor input, int64_t num_groups, std::optional<Tensor> weight, std::optional<Tensor> bias, double eps, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor group_norm(Tensor input, int64_t num_groups, std::optional<Tensor> weight, std::optional<Tensor> bias, double eps);
+void group_norm_(Tensor input, int64_t num_groups, std::optional<Tensor> weight, std::optional<Tensor> bias, double eps, Tensor output);
+
+} // namespace infinicore::op

include/infinicore/ops/isclose.hpp

@@ -0,0 +1,16 @@
+#pragma once
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class IsClose {
+public:
+    using schema = void (*)(Tensor, Tensor, double, double, bool, Tensor);
+    static void execute(Tensor a, Tensor b, double rtol, double atol, bool equal_nan, Tensor out);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor isclose(Tensor a, Tensor b, double rtol = 1e-05, double atol = 1e-08, bool equal_nan = false);
+
+} // namespace infinicore::op

include/infinicore/ops/reshape.hpp

@@ -0,0 +1,10 @@
+#pragma once
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+// Reshape 主要是元数据操作，不需要像 Compute Op 那样定义 Schema/Dispatcher
+Tensor reshape(Tensor input, Shape shape);
+
+} // namespace infinicore::op

include/infinicore/ops/rrelu.hpp

@@ -0,0 +1,18 @@
+#pragma once
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class RReLU {
+public:
+    using schema = void (*)(Tensor, double, double, bool, Tensor);
+    static void execute(Tensor input, double lower, double upper, bool training, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor rrelu(Tensor input, double lower = 0.125, double upper = 0.3333333333333333, bool training = false, bool inplace = false);
+// In-place
+Tensor& rrelu_(Tensor input, double lower = 0.125, double upper = 0.3333333333333333, bool training = false);
+
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -54,6 +54,9 @@
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.unsqueeze import unsqueeze
+from infinicore.ops.reshape import reshape
+from infinicore.ops.bitwise_xor import bitwise_xor
+from infinicore.ops.isclose import isclose
 from infinicore.tensor import (
     Tensor,
     empty,
@@ -134,6 +137,9 @@
     "strided_empty",
     "strided_from_blob",
     "zeros",
+    "reshape",
+    "bitwise_xor",
+    "isclose",
 ]
 
 use_ntops = False

python/infinicore/nn/functional/__init__.py

@@ -6,6 +6,8 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
+from .group_norm import group_norm
+from .rrelu import rrelu
 
 __all__ = [
     "causal_softmax",
@@ -17,4 +19,6 @@
     "embedding",
     "rope",
     "RopeAlgo",
+    "group_norm",
+    "rrelu",
 ]

python/infinicore/nn/functional/group_norm.py

@@ -0,0 +1,32 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def group_norm(
+    input: Tensor,
+    num_groups: int,
+    weight: Tensor | None = None,
+    bias: Tensor | None = None,
+    eps: float = 1e-5,
+) -> Tensor:
+    r"""Applies Group Normalization."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.group_norm(
+            input, num_groups, weight=weight, bias=bias, eps=eps
+        )
+
+    weight_underlying = weight._underlying if weight is not None else None
+    bias_underlying = bias._underlying if bias is not None else None
+
+    # Call C++ Backend via PyBind
+    return Tensor(
+        _infinicore.group_norm(
+            input._underlying, 
+            num_groups, 
+            weight_underlying, 
+            bias_underlying, 
+            eps
+        )
+    )

python/infinicore/nn/functional/rrelu.py

@@ -0,0 +1,11 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def rrelu(input: Tensor, lower: float = 0.125, upper: float = 0.3333333333333333, training: bool = False, inplace: bool = False, generator=None) -> Tensor:
+    r"""Applies the randomized leaky rectified linear unit function."""
+    
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.rrelu(input, lower, upper, training, inplace, generator)
+    
+    return Tensor(_infinicore.rrelu(input._underlying, lower, upper, training, inplace))

python/infinicore/ops/bitwise_xor.py

@@ -0,0 +1,15 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def bitwise_xor(input: Tensor, other: Tensor, out: Tensor | None = None) -> Tensor:
+    """Computes the bitwise XOR of input and other."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.bitwise_xor(input, other, out=out)
+    
+    if out is not None:
+        _infinicore.bitwise_xor_out(input._underlying, other._underlying, out._underlying)
+        return out
+    else:
+        return Tensor(_infinicore.bitwise_xor(input._underlying, other._underlying))