Add InfiniCore axpy, blas_copy, scal and swap wrappers

Author: xu-zhengzhong

include/infinicore/ops.hpp

@@ -10,11 +10,13 @@
 #include "ops/atanh.hpp"
 #include "ops/attention.hpp"
 #include "ops/avg_pool1d.hpp"
+#include "ops/axpy.hpp"
 #include "ops/baddbmm.hpp"
 #include "ops/bilinear.hpp"
 #include "ops/binary_cross_entropy_with_logits.hpp"
 #include "ops/blas_amax.hpp"
 #include "ops/blas_amin.hpp"
+#include "ops/blas_copy.hpp"
 #include "ops/blas_dot.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/cdist.hpp"
@@ -46,8 +48,10 @@
 #include "ops/relu.hpp"
 #include "ops/rms_norm.hpp"
 #include "ops/rope.hpp"
+#include "ops/scal.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"
 #include "ops/softmax.hpp"
+#include "ops/swap.hpp"
 #include "ops/swiglu.hpp"
 #include "ops/topksoftmax.hpp"

include/infinicore/ops/axpy.hpp

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

include/infinicore/ops/blas_copy.hpp

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

include/infinicore/ops/scal.hpp

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

include/infinicore/ops/swap.hpp

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

python/infinicore/__init__.py

@@ -61,10 +61,12 @@
 from infinicore.ops.asum import asum
 from infinicore.ops.atanh import atanh
 from infinicore.ops.attention import attention
+from infinicore.ops.axpy import axpy
 from infinicore.ops.baddbmm import baddbmm
 from infinicore.ops.bilinear import bilinear
 from infinicore.ops.blas_amax import blas_amax
 from infinicore.ops.blas_amin import blas_amin
+from infinicore.ops.blas_copy import blas_copy
 from infinicore.ops.blas_dot import blas_dot
 from infinicore.ops.binary_cross_entropy_with_logits import (
     binary_cross_entropy_with_logits,
@@ -112,8 +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.scal import scal
 from infinicore.ops.scatter import scatter
 from infinicore.ops.sinh import sinh
+from infinicore.ops.swap import swap
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.sum import sum
 from infinicore.ops.take import take
@@ -191,8 +195,10 @@
     "argwhere",
     "asin",
     "asum",
+    "axpy",
     "blas_amax",
     "blas_amin",
+    "blas_copy",
     "blas_dot",
     "acos",
     "addbmm",
@@ -241,6 +247,7 @@
     "float_power",
     "flipud",
     "scatter",
+    "scal",
     "logcumsumexp",
     "logical_not",
     "logical_and",
@@ -253,6 +260,7 @@
     "index_add",
     "take",
     "sinh",
+    "swap",
     "ones",
     "broadcast_to",
     "strided_empty",

python/infinicore/ops/axpy.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def axpy(alpha: Tensor, x: Tensor, y: Tensor):
+    _infinicore.axpy_(alpha._underlying, x._underlying, y._underlying)
+    return y

python/infinicore/ops/blas_copy.py

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

python/infinicore/ops/scal.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def scal(x: Tensor, alpha: Tensor):
+    _infinicore.scal_(x._underlying, alpha._underlying)
+    return x

python/infinicore/ops/swap.py

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