[2025秋季][T1-1-4] Gary828

include/infinicore/ops/all.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+#include <vector>
+namespace infinicore::op {
+class All {
+public:
+    using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool);
+    static void execute(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor all(Tensor input, std::vector<size_t> dim, bool keepdim = false);
+void all_(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
+
+} // namespace infinicore::op

include/infinicore/ops/sum.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+#include <vector>
+
+namespace infinicore::op {
+class Sum {
+public:
+    using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool);
+    static void execute(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor sum(Tensor input, std::vector<size_t> dim, bool keepdim = false);
+void sum_(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
+
+} // namespace infinicore::op

include/infinicore/ops/topk.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+namespace infinicore::op {
+class TopK {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, size_t, size_t, bool, bool);
+    static void execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::pair<Tensor, Tensor> topk(Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
+void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
+
+} // namespace infinicore::op

include/infinicore/ops/var.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+#include <utility>
+#include <vector>
+namespace infinicore::op {
+class Var {
+public:
+    using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool, bool); // var_output, input, dim, unbiased, keepdim
+    static void execute(Tensor var_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor var(Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+void var_(Tensor var_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+
+} // namespace infinicore::op

include/infinicore/ops/var_mean.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+#include <utility>
+#include <vector>
+namespace infinicore::op {
+class Var_Mean {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, std::vector<size_t>, bool, bool); // var_output, mean_output, input, dim, unbiased, keepdim
+    static void execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::pair<Tensor, Tensor> var_mean(Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -3,6 +3,7 @@
 
 #include "infiniop/handle.h"
 #include "infiniop/ops/add.h"
+#include "infiniop/ops/all.h"
 #include "infiniop/ops/attention.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"
@@ -28,10 +29,14 @@
 #include "infiniop/ops/softmax.h"
 #include "infiniop/ops/softplus.h"
 #include "infiniop/ops/sub.h"
+#include "infiniop/ops/sum.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/ops/tanh.h"
+#include "infiniop/ops/topk.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
+#include "infiniop/ops/var.h"
+#include "infiniop/ops/var_mean.h"
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
 

include/infiniop/ops/all.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_ALL_API_H__
+#define __INFINIOP_ALL_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstddef>
+#include <vector>
+typedef struct InfiniopDescriptor *infiniopAllDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateAllDescriptor(infiniopHandle_t handle,
+                                                        infiniopAllDescriptor_t *desc_ptr,
+                                                        infiniopTensorDescriptor_t output_desc,
+                                                        infiniopTensorDescriptor_t input_desc,
+                                                        size_t *dim,
+                                                        size_t dim_size,
+                                                        bool keepdim);
+
+__C __export infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopAll(infiniopAllDescriptor_t desc,
+                                        void *workspace,
+                                        size_t workspace_size,
+                                        void *output,
+                                        const void *input,
+                                        size_t *dim,
+                                        size_t dim_size,
+                                        bool keepdim,
+                                        void *stream);
+
+__C __export infiniStatus_t infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc);
+
+#endif

include/infiniop/ops/sum.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_SUM_API_H__
+#define __INFINIOP_SUM_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstddef>
+#include <vector>
+typedef struct InfiniopDescriptor *infiniopSumDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateSumDescriptor(infiniopHandle_t handle,
+                                                        infiniopSumDescriptor_t *desc_ptr,
+                                                        infiniopTensorDescriptor_t output_desc,
+                                                        infiniopTensorDescriptor_t input_desc,
+                                                        size_t *dim,
+                                                        size_t dim_size,
+                                                        bool keepdim);
+
+__C __export infiniStatus_t infiniopGetSumWorkspaceSize(infiniopSumDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopSum(infiniopSumDescriptor_t desc,
+                                        void *workspace,
+                                        size_t workspace_size,
+                                        void *output,
+                                        const void *input,
+                                        size_t *dim,
+                                        size_t dim_size,
+                                        bool keepdim,
+                                        void *stream);
+
+__C __export infiniStatus_t infiniopDestroySumDescriptor(infiniopSumDescriptor_t desc);
+
+#endif

include/infiniop/ops/topk.h

@@ -0,0 +1,35 @@
+#ifndef __INFINIOP_TOPK_API_H__
+#define __INFINIOP_TOPK_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstddef>
+#include <vector>
+typedef struct InfiniopDescriptor *infiniopTopKDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTopKDescriptor(infiniopHandle_t handle,
+                                                         infiniopTopKDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t values_output_desc,
+                                                         infiniopTensorDescriptor_t indices_output_desc,
+                                                         infiniopTensorDescriptor_t input_desc,
+                                                         size_t k,
+                                                         size_t dim,
+                                                         bool largest,
+                                                         bool sorted);
+
+__C __export infiniStatus_t infiniopGetTopKWorkspaceSize(infiniopTopKDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTopK(infiniopTopKDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *values_output,
+                                         void *indices_output,
+                                         const void *input,
+                                         size_t k,
+                                         size_t dim,
+                                         bool largest,
+                                         bool sorted,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTopKDescriptor(infiniopTopKDescriptor_t desc);
+
+#endif

include/infiniop/ops/var.h

@@ -0,0 +1,33 @@
+#ifndef __INFINIOP_VAR_API_H__
+#define __INFINIOP_VAR_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstddef>
+#include <vector>
+typedef struct InfiniopDescriptor *infiniopVarDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateVarDescriptor(infiniopHandle_t handle,
+                                                        infiniopVarDescriptor_t *desc_ptr,
+                                                        infiniopTensorDescriptor_t var_output_desc,
+                                                        infiniopTensorDescriptor_t input_desc,
+                                                        size_t *dim,
+                                                        size_t dim_size,
+                                                        bool unbiased,
+                                                        bool keepdim);
+
+__C __export infiniStatus_t infiniopGetVarWorkspaceSize(infiniopVarDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopVar(infiniopVarDescriptor_t desc,
+                                        void *workspace,
+                                        size_t workspace_size,
+                                        void *var_output,
+                                        const void *input,
+                                        size_t *dim,
+                                        size_t dim_size,
+                                        bool unbiased,
+                                        bool keepdim,
+                                        void *stream);
+
+__C __export infiniStatus_t infiniopDestroyVarDescriptor(infiniopVarDescriptor_t desc);
+
+#endif

include/infiniop/ops/var_mean.h

@@ -0,0 +1,35 @@
+#ifndef __INFINIOP_VAR_MEAN_API_H__
+#define __INFINIOP_VAR_MEAN_API_H__
+
+#include "../operator_descriptor.h"
+#include <cstddef>
+#include <vector>
+typedef struct InfiniopDescriptor *infiniopVarMeanDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateVarMeanDescriptor(infiniopHandle_t handle,
+                                                            infiniopVarMeanDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t var_output_desc,
+                                                            infiniopTensorDescriptor_t mean_output_desc,
+                                                            infiniopTensorDescriptor_t input_desc,
+                                                            size_t *dim,
+                                                            size_t dim_size,
+                                                            bool unbiased,
+                                                            bool keepdim);
+
+__C __export infiniStatus_t infiniopGetVarMeanWorkspaceSize(infiniopVarMeanDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopVarMean(infiniopVarMeanDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *var_output,
+                                            void *mean_output,
+                                            const void *input,
+                                            size_t *dim,
+                                            size_t dim_size,
+                                            bool unbiased,
+                                            bool keepdim,
+                                            void *stream);
+
+__C __export infiniStatus_t infiniopDestroyVarMeanDescriptor(infiniopVarMeanDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -48,6 +48,11 @@
 from infinicore.ops.paged_attention_prefill import paged_attention_prefill
 from infinicore.ops.paged_caching import paged_caching
 from infinicore.ops.rearrange import rearrange
+from infinicore.ops.sum import sum
+from infinicore.ops.var_mean import var_mean
+from infinicore.ops.var import var
+from infinicore.ops.topk import topk
+from infinicore.ops.all import all
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.unsqueeze import unsqueeze
 from infinicore.tensor import (
@@ -125,6 +130,11 @@
     "strided_empty",
     "strided_from_blob",
     "zeros",
+    "sum",
+    "var_mean",
+    "var",
+    "topk",
+    "all",
 ]
 
 use_ntops = False

python/infinicore/ops/all.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def all(input, dim=None, keepdim=False, out=None):
+    if out is None:
+        return Tensor(_infinicore.all(input._underlying, dim, keepdim))
+
+    _infinicore.all_(out._underlying, input._underlying, dim, keepdim)
+
+    return out

python/infinicore/ops/sum.py

@@ -0,0 +1,28 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def sum(input, dim=None, keepdim=False, out=None):
+    """
+    Sum the elements of the input tensor along the given dimensions.
+
+    Args:
+        input (Tensor): The input tensor.
+        out (Tensor, optional): The output tensor.
+
+    Returns:
+        Tensor: The output tensor.
+
+    Example:
+        >>> import infinicore
+        >>> input = infinicore.tensor([[1, 2, 3], [4, 5, 6]])
+        >>> output = infinicore.sum(input)
+        >>> print(output)
+        tensor([15])
+    """
+    if out is None:
+        return Tensor(_infinicore.sum(input._underlying, dim, keepdim))
+
+    _infinicore.sum_(out._underlying, input._underlying, dim, keepdim)
+
+    return out

python/infinicore/ops/topk.py

@@ -0,0 +1,12 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def topk(input, k, dim, largest=True, sorted=True, out=None):
+    if out is None:
+        values, indices = _infinicore.topk(input._underlying, k, dim, largest, sorted)
+        return Tensor(values), Tensor(indices)
+
+    _infinicore.topk_(out._underlying, input._underlying, k, dim, largest, sorted)
+
+    return out

python/infinicore/ops/var.py

@@ -0,0 +1,12 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def var(input, dim=None, unbiased=True, keepdim=False, out=None):
+    if out is None:
+        var_tensor = _infinicore.var(input._underlying, dim, unbiased, keepdim)
+        return Tensor(var_tensor)
+    var_output = out
+    _infinicore.var_(var_output._underlying, input._underlying, dim, unbiased, keepdim)
+
+    return out