Add InfiniCore rot, rotg, rotm and rotmg wrappers

Author: xu-zhengzhong

include/infinicore/ops.hpp

@@ -48,6 +48,10 @@
 #include "ops/relu.hpp"
 #include "ops/rms_norm.hpp"
 #include "ops/rope.hpp"
+#include "ops/rot.hpp"
+#include "ops/rotg.hpp"
+#include "ops/rotm.hpp"
+#include "ops/rotmg.hpp"
 #include "ops/scal.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"

include/infinicore/ops/rot.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Rot {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor);
+    static void execute(Tensor x, Tensor y, Tensor c, Tensor s);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+void rot_(Tensor x, Tensor y, Tensor c, Tensor s);
+
+} // namespace infinicore::op

include/infinicore/ops/rotg.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Rotg {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor);
+    static void execute(Tensor x, Tensor y, Tensor c, Tensor s);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+void rotg_(Tensor x, Tensor y, Tensor c, Tensor s);
+
+} // namespace infinicore::op

include/infinicore/ops/rotm.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Rotm {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+    static void execute(Tensor x, Tensor y, Tensor param);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+void rotm_(Tensor x, Tensor y, Tensor param);
+
+} // namespace infinicore::op

include/infinicore/ops/rotmg.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Rotmg {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor);
+    static void execute(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+void rotmg_(Tensor d1, Tensor d2, Tensor x1, Tensor y1, Tensor param);
+
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -114,6 +114,10 @@
 from infinicore.ops.paged_caching import paged_caching
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.reciprocal import reciprocal
+from infinicore.ops.rot import rot
+from infinicore.ops.rotg import rotg
+from infinicore.ops.rotm import rotm
+from infinicore.ops.rotmg import rotmg
 from infinicore.ops.scal import scal
 from infinicore.ops.scatter import scatter
 from infinicore.ops.sinh import sinh
@@ -247,6 +251,10 @@
     "float_power",
     "flipud",
     "scatter",
+    "rot",
+    "rotg",
+    "rotm",
+    "rotmg",
     "scal",
     "logcumsumexp",
     "logical_not",

python/infinicore/ops/rot.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def rot(x: Tensor, y: Tensor, c: Tensor, s: Tensor):
+    _infinicore.rot_(x._underlying, y._underlying, c._underlying, s._underlying)
+    return x, y

python/infinicore/ops/rotg.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def rotg(x: Tensor, y: Tensor, c: Tensor, s: Tensor):
+    _infinicore.rotg_(x._underlying, y._underlying, c._underlying, s._underlying)
+    return x, y, c, s

python/infinicore/ops/rotm.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def rotm(x: Tensor, y: Tensor, param: Tensor):
+    _infinicore.rotm_(x._underlying, y._underlying, param._underlying)
+    return x, y

python/infinicore/ops/rotmg.py

@@ -0,0 +1,13 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def rotmg(d1: Tensor, d2: Tensor, x1: Tensor, y1: Tensor, param: Tensor):
+    _infinicore.rotmg_(
+        d1._underlying,
+        d2._underlying,
+        x1._underlying,
+        y1._underlying,
+        param._underlying,
+    )
+    return d1, d2, x1, param