Finish T1-1-4: sum、topk、var、var_mean、all

Author: PPPoint-t

include/infinicore/ops/all.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class AllGlobal {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor input, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor all_global(Tensor input);
+void all_global_(Tensor input, Tensor output);
+
+class AllReduce {
+public:
+    using schema = void (*)(Tensor, Tensor, int, bool);
+    static void execute(Tensor input, Tensor output, int dim, bool keepdim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor all_reduce(Tensor input, int dim, bool keepdim);
+void all_reduce_(Tensor input, Tensor output, int dim, bool keepdim);
+
+} // namespace infinicore::op

include/infinicore/ops/sum.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class SumGlobal {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor input, Tensor output);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor sum_global(Tensor input);
+void sum_global_(Tensor input, Tensor output);
+
+class SumReduce {
+public:
+    using schema = void (*)(Tensor, Tensor, int, bool);
+    static void execute(Tensor input, Tensor output, int dim, bool keepdim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor sum_reduce(Tensor input, int dim, bool keepdim);
+void sum_reduce_(Tensor input, Tensor output, int dim, bool keepdim);
+
+} // namespace infinicore::op

include/infinicore/ops/topk.hpp

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

include/infinicore/ops/var.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class VarGlobal {
+public:
+    using schema = void (*)(Tensor, Tensor, int);
+    static void execute(Tensor input, Tensor output, int correction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor var_global(Tensor input, int correction);
+void var_global_(Tensor input, Tensor output, int correction);
+
+class VarReduce {
+public:
+    using schema = void (*)(Tensor, Tensor, int, int, bool);
+    static void execute(Tensor input, Tensor output, int dim, int correction, bool keepdim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor var_reduce(Tensor input, int dim, int correction, bool keepdim);
+void var_reduce_(Tensor input, Tensor output, int dim, int correction, bool keepdim);
+
+} // namespace infinicore::op

include/infinicore/ops/var_mean.hpp

@@ -0,0 +1,29 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <tuple>
+
+namespace infinicore::op {
+
+class VarMeanGlobal {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, int);
+    static void execute(Tensor input, Tensor out_var, Tensor out_mean, int correction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::tuple<Tensor, Tensor> var_mean_global(Tensor input, int correction);
+void var_mean_global_(Tensor input, Tensor out_var, Tensor out_mean, int correction);
+
+class VarMeanReduce {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, int, int, bool);
+    static void execute(Tensor input, Tensor out_var, Tensor out_mean, int dim, int correction, bool keepdim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+std::tuple<Tensor, Tensor> var_mean_reduce(Tensor input, int dim, int correction, bool keepdim);
+void var_mean_reduce_(Tensor input, Tensor out_var, Tensor out_mean, int dim, int correction, bool keepdim);
+
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -40,13 +40,18 @@
     uint8,
 )
 from infinicore.ops.add import add
+from infinicore.ops.all import all
 from infinicore.ops.attention import attention
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.squeeze import squeeze
+from infinicore.ops.sum import sum
+from infinicore.ops.topk import topk
 from infinicore.ops.unsqueeze import unsqueeze
+from infinicore.ops.var import var
+from infinicore.ops.var_mean import var_mean
 from infinicore.tensor import (
     Tensor,
     empty,
@@ -119,6 +124,11 @@
     "strided_empty",
     "strided_from_blob",
     "zeros",
+    "sum",
+    "topk",
+    "var",
+    "var_mean",
+    "all",
 ]
 
 use_ntops = False

python/infinicore/ops/all.py

@@ -0,0 +1,57 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def all(
+    input: Tensor,
+    dim: int | tuple[int] | list[int] | None = None,
+    keepdim: bool = False,
+    *,
+    out=None,
+) -> Tensor:
+    r"""Computes the logical AND of all elements."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.all(input, dim=dim, keepdim=keepdim, out=out)
+
+    if dim is None:
+        if out is None:
+            return Tensor(_infinicore.all_global(input._underlying))
+        _infinicore.all_global_(input._underlying, out._underlying)
+        return out
+
+    else:
+        if isinstance(dim, int):
+            dims = [dim]
+        else:
+            dims = list(dim)
+
+        ndim = input.ndim
+        normalized_dims = sorted(
+            [d if d >= 0 else d + ndim for d in dims], reverse=True
+        )
+
+        current_input = input
+
+        if len(normalized_dims) == 1 and out is not None:
+            _infinicore.all_reduce_(
+                current_input._underlying, out._underlying, normalized_dims[0], keepdim
+            )
+            return out
+
+        for i, target_dim in enumerate(normalized_dims):
+            is_last_step = i == len(normalized_dims) - 1
+
+            if is_last_step and out is not None:
+                _infinicore.all_reduce_(
+                    current_input._underlying, out._underlying, target_dim, keepdim
+                )
+                return out
+            else:
+                res_ptr = _infinicore.all_reduce(
+                    current_input._underlying, target_dim, keepdim
+                )
+                current_input = Tensor(res_ptr)
+
+        return current_input

python/infinicore/ops/sum.py

@@ -0,0 +1,38 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def sum(
+    input: Tensor,
+    dim: int | tuple[int] | list[int] | None = None,
+    keepdim=False,
+    *,
+    dtype=None,
+    out=None,
+) -> Tensor:
+    r"""Apply the sum function."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.sum(
+            input, dim, keepdim=keepdim, dtype=dtype, out=out
+        )
+
+    if dim is None:
+        if out is None:
+            return Tensor(_infinicore.sum_global(input._underlying))
+        _infinicore.sum_global_(input._underlying, out._underlying)
+        return out
+
+    else:
+        target_dim = dim
+        if isinstance(target_dim, (tuple, list)):
+            if len(target_dim) == 1:
+                target_dim = target_dim[0]
+        if out is None:
+            return Tensor(
+                _infinicore.sum_reduce(input._underlying, target_dim, keepdim)
+            )
+
+        _infinicore.sum_reduce_(input._underlying, out._underlying, target_dim, keepdim)
+        return out

python/infinicore/ops/topk.py

@@ -0,0 +1,39 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def topk(
+    input: Tensor,
+    k: int,
+    dim: int = -1,
+    largest: bool = True,
+    sorted: bool = True,
+    *,
+    out=None,
+):
+    r"""Returns the k largest elements of the given input tensor along a given dimension."""
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.topk(input, k, dim, largest, sorted, out=out)
+
+    if out is None:
+        res_values, res_indices = _infinicore.topk(
+            input._underlying, k, dim, largest, sorted
+        )
+        return Tensor(res_values), Tensor(res_indices)
+    else:
+        if not isinstance(out, (tuple, list)) or len(out) != 2:
+            raise ValueError("out argument must be a tuple of (values, indices)")
+
+        out_values, out_indices = out
+        _infinicore.topk_(
+            input._underlying,
+            out_values._underlying,
+            out_indices._underlying,
+            k,
+            dim,
+            largest,
+            sorted,
+        )
+        return out_values, out_indices

python/infinicore/ops/var.py

@@ -0,0 +1,58 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def var(
+    input: Tensor,
+    dim: int | tuple[int] | list[int] | None = None,
+    unbiased: bool | None = None,
+    correction: int | None = None,
+    keepdim: bool = False,
+    *,
+    dtype=None,
+    out=None,
+) -> Tensor:
+    r"""Returns the variance of the input tensor."""
+
+    if unbiased is not None:
+        if correction is not None and correction != (1 if unbiased else 0):
+            raise ValueError(
+                "Cannot specify both 'unbiased' and 'correction' with conflicting values."
+            )
+        final_correction = 1 if unbiased else 0
+    else:
+        final_correction = correction if correction is not None else 1
+
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.var(
+            input,
+            dim=dim,
+            correction=final_correction,
+            keepdim=keepdim,
+            dtype=dtype,
+            out=out,
+        )
+
+    if dim is None:
+        if out is None:
+            return Tensor(_infinicore.var_global(input._underlying, final_correction))
+        _infinicore.var_global_(input._underlying, out._underlying, final_correction)
+        return out
+    else:
+        target_dim = dim
+        if isinstance(target_dim, (tuple, list)):
+            if len(target_dim) == 1:
+                target_dim = target_dim[0]
+
+        if out is None:
+            return Tensor(
+                _infinicore.var_reduce(
+                    input._underlying, target_dim, final_correction, keepdim
+                )
+            )
+
+        _infinicore.var_reduce_(
+            input._underlying, out._underlying, target_dim, final_correction, keepdim
+        )
+        return out