Add InfiniCore asum, blas_amax, blas_amin, blas_dot and nrm2 wrappers

Author: xu-zhengzhong

include/infinicore/ops.hpp

@@ -6,12 +6,16 @@
 #include "ops/addcmul.hpp"
 #include "ops/asin.hpp"
 #include "ops/asinh.hpp"
+#include "ops/asum.hpp"
 #include "ops/atanh.hpp"
 #include "ops/attention.hpp"
 #include "ops/avg_pool1d.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_dot.hpp"
 #include "ops/causal_softmax.hpp"
 #include "ops/cdist.hpp"
 #include "ops/conv2d.hpp"
@@ -28,6 +32,7 @@
 #include "ops/layer_norm.hpp"
 #include "ops/linear.hpp"
 #include "ops/matmul.hpp"
+#include "ops/nrm2.hpp"
 #include "ops/ones.hpp"
 #include "ops/paged_attention.hpp"
 #include "ops/paged_attention_prefill.hpp"

include/infinicore/ops/asum.hpp

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

include/infinicore/ops/blas_amax.hpp

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

include/infinicore/ops/blas_amin.hpp

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

include/infinicore/ops/blas_dot.hpp

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

include/infinicore/ops/nrm2.hpp

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

python/infinicore/__init__.py

@@ -58,10 +58,14 @@
 from infinicore.ops.argwhere import argwhere
 from infinicore.ops.asin import asin
 from infinicore.ops.asinh import asinh
+from infinicore.ops.asum import asum
 from infinicore.ops.atanh import atanh
 from infinicore.ops.attention import attention
 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_dot import blas_dot
 from infinicore.ops.binary_cross_entropy_with_logits import (
     binary_cross_entropy_with_logits,
 )
@@ -102,6 +106,7 @@
 from infinicore.ops.mha_varlen import mha_varlen
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
+from infinicore.ops.nrm2 import nrm2
 from infinicore.ops.paged_attention import paged_attention
 from infinicore.ops.paged_attention_prefill import paged_attention_prefill
 from infinicore.ops.paged_caching import paged_caching
@@ -185,6 +190,10 @@
     "add_rms_norm",
     "argwhere",
     "asin",
+    "asum",
+    "blas_amax",
+    "blas_amin",
+    "blas_dot",
     "acos",
     "addbmm",
     "floor",
@@ -210,6 +219,7 @@
     "dist",
     "logdet",
     "narrow",
+    "nrm2",
     "ldexp",
     "lerp",
     "kthvalue",

python/infinicore/ops/asum.py

@@ -0,0 +1,10 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def asum(x: Tensor, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.asum(x._underlying))
+
+    _infinicore.asum_(out._underlying, x._underlying)
+    return out

python/infinicore/ops/blas_amax.py

@@ -0,0 +1,10 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def blas_amax(x: Tensor, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.blas_amax(x._underlying))
+
+    _infinicore.blas_amax_(out._underlying, x._underlying)
+    return out

python/infinicore/ops/blas_amin.py

@@ -0,0 +1,10 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def blas_amin(x: Tensor, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.blas_amin(x._underlying))
+
+    _infinicore.blas_amin_(out._underlying, x._underlying)
+    return out