[2025秋季][T1-1-1]guozhihao-224

include/infinicore/ops/logical_or.hpp

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

include/infinicore/ops/logical_xor.hpp

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

include/infinicore/ops/logsigmoid.hpp

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

include/infinicore/ops/vdot.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Vdot {
+public:
+    using schema = void (*)(Tensor out, Tensor a, Tensor b);
+    static void execute(Tensor out, Tensor a, Tensor b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor vdot(Tensor a, Tensor b);
+void vdot_(Tensor out, Tensor a, Tensor b);
+
+} // namespace infinicore::op
+
+

include/infinicore/ops/where.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Where {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor);
+    static void execute(Tensor out, Tensor cond, Tensor x, Tensor y);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor where(Tensor cond, Tensor x, Tensor y);
+void where_(Tensor out, Tensor cond, Tensor x, Tensor y);
+
+class WhereIndices {
+public:
+    using schema = std::vector<Tensor> (*)(Tensor);
+    static std::vector<Tensor> execute(Tensor cond);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// where(cond) -> tuple of index tensors
+std::vector<Tensor> where_indices(Tensor cond);
+
+} // namespace infinicore::op
+
+

include/infiniop.h

@@ -11,6 +11,9 @@
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
 #include "infiniop/ops/layer_norm.h"
+#include "infiniop/ops/logical_or.h"
+#include "infiniop/ops/logical_xor.h"
+#include "infiniop/ops/logsigmoid.h"
 #include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/lp_norm.h"
 #include "infiniop/ops/mul.h"
@@ -29,6 +32,8 @@
 #include "infiniop/ops/tanh.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
+#include "infiniop/ops/where.h"
+#include "infiniop/ops/vdot.h"
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
 

include/infiniop/ops/logical_or.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_LOGICAL_OR_API_H__
+#define __INFINIOP_LOGICAL_OR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogicalOrDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogicalOrDescriptor(infiniopHandle_t handle,
+                                                                 infiniopLogicalOrDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t c_desc,
+                                                                 infiniopTensorDescriptor_t a_desc,
+                                                                 infiniopTensorDescriptor_t b_desc);
+
+__C __export infiniStatus_t infiniopGetLogicalOrWorkspaceSize(infiniopLogicalOrDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogicalOr(infiniopLogicalOrDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *c,
+                                               const void *a,
+                                               const void *b,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogicalOrDescriptor(infiniopLogicalOrDescriptor_t desc);
+
+#endif
+

include/infiniop/ops/logical_xor.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_LOGICAL_XOR_API_H__
+#define __INFINIOP_LOGICAL_XOR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogicalXorDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogicalXorDescriptor(infiniopHandle_t handle,
+                                                                 infiniopLogicalXorDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t c_desc,
+                                                                 infiniopTensorDescriptor_t a_desc,
+                                                                 infiniopTensorDescriptor_t b_desc);
+
+__C __export infiniStatus_t infiniopGetLogicalXorWorkspaceSize(infiniopLogicalXorDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogicalXor(infiniopLogicalXorDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *c,
+                                               const void *a,
+                                               const void *b,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogicalXorDescriptor(infiniopLogicalXorDescriptor_t desc);
+
+#endif
+

include/infiniop/ops/logsigmoid.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_LOGSIGMOID_API_H__
+#define __INFINIOP_LOGSIGMOID_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogSigmoidDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogSigmoidDescriptor(infiniopHandle_t handle,
+                                                                 infiniopLogSigmoidDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t y,
+                                                                 infiniopTensorDescriptor_t x);
+
+__C __export infiniStatus_t infiniopGetLogSigmoidWorkspaceSize(infiniopLogSigmoidDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogSigmoid(infiniopLogSigmoidDescriptor_t desc,
+                                                void *workspace,
+                                                size_t workspace_size,
+                                                void *y,
+                                                const void *x,
+                                                void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogSigmoidDescriptor(infiniopLogSigmoidDescriptor_t desc);
+
+#endif
+

include/infiniop/ops/vdot.h

@@ -0,0 +1,34 @@
+#ifndef __INFINIOP_VDOT_API_H__
+#define __INFINIOP_VDOT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopVdotDescriptor_t;
+
+// out = vdot(a, b)
+__C __export infiniStatus_t infiniopCreateVdotDescriptor(
+    infiniopHandle_t handle,
+    infiniopVdotDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t a_desc,
+    infiniopTensorDescriptor_t b_desc);
+
+__C __export infiniStatus_t infiniopGetVdotWorkspaceSize(
+    infiniopVdotDescriptor_t desc,
+    size_t *size);
+
+__C __export infiniStatus_t infiniopVdot(
+    infiniopVdotDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *a,
+    const void *b,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyVdotDescriptor(
+    infiniopVdotDescriptor_t desc);
+
+#endif // __INFINIOP_VDOT_API_H__
+
+

include/infiniop/ops/where.h

@@ -0,0 +1,60 @@
+#ifndef __INFINIOP_WHERE_API_H__
+#define __INFINIOP_WHERE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopWhereDescriptor_t;
+
+// y = where(cond, x, y)
+__C __export infiniStatus_t infiniopCreateWhereDescriptor(
+    infiniopHandle_t handle,
+    infiniopWhereDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t cond_desc,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t y_desc);
+
+__C __export infiniStatus_t infiniopGetWhereWorkspaceSize(
+    infiniopWhereDescriptor_t desc,
+    size_t *size);
+
+__C __export infiniStatus_t infiniopWhere(
+    infiniopWhereDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *cond,
+    const void *x,
+    const void *y,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyWhereDescriptor(
+    infiniopWhereDescriptor_t desc);
+
+// where(cond) -> indices tuple
+typedef struct InfiniopDescriptor *infiniopWhereIndicesDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateWhereIndicesDescriptor(
+    infiniopHandle_t handle,
+    infiniopWhereIndicesDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t cond_desc);
+
+__C __export infiniStatus_t infiniopGetWhereIndicesWorkspaceSize(
+    infiniopWhereIndicesDescriptor_t desc,
+    size_t *size);
+
+__C __export infiniStatus_t infiniopWhereIndices(
+    infiniopWhereIndicesDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void **outputs,  // NDIM 个输出张量的指针数组
+    const void *cond,
+    void *stream,
+    size_t *num_true);  // 输出：True 元素的数量
+
+__C __export infiniStatus_t infiniopDestroyWhereIndicesDescriptor(
+    infiniopWhereIndicesDescriptor_t desc);
+
+#endif
+
+

python/infinicore/__init__.py

@@ -41,6 +41,11 @@
 )
 from infinicore.ops.add import add
 from infinicore.ops.attention import attention
+from infinicore.ops.logical_or import logical_or
+from infinicore.ops.logical_xor import logical_xor
+from infinicore.ops.logsigmoid import logsigmoid
+from infinicore.ops.where import where
+from infinicore.ops.vdot import vdot
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
@@ -101,6 +106,11 @@
     # Operations.
     "add",
     "attention",
+    "logical_or",
+    "logical_xor",
+    "logsigmoid",
+    "where",
+    "vdot",
     "matmul",
     "mul",
     "narrow",

python/infinicore/nn/functional/__init__.py

@@ -1,6 +1,7 @@
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .linear import linear
+from .logsigmoid import logsigmoid
 from .random_sample import random_sample
 from .rms_norm import rms_norm
 from .rope import RopeAlgo, rope
@@ -17,4 +18,5 @@
     "embedding",
     "rope",
     "RopeAlgo",
+    "logsigmoid",
 ]

python/infinicore/nn/functional/logsigmoid.py

@@ -0,0 +1,14 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def logsigmoid(input: Tensor, out=None) -> Tensor:
+    """Apply elementwise log-sigmoid."""
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.logsigmoid(input, out=out)
+    if out is None:
+        return Tensor(_infinicore.logsigmoid(input._underlying))
+
+    _infinicore.logsigmoid_(out._underlying, input._underlying)
+    return out

python/infinicore/nn/functional/silu.py

@@ -6,7 +6,7 @@
 def silu(input: Tensor, inplace: bool = False, *, out=None) -> Tensor:
     r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise."""
 
-    if infinicore.use_ntops and input.device.type in ("cuda", "musa") and out is None:
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
         return infinicore.ntops.torch.silu(input, inplace=inplace)
 
     if inplace:

python/infinicore/ops/logical_or.py

@@ -0,0 +1,15 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def logical_or(input, other, *, out=None):
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.logical_or(input, other, out=out)
+
+    if out is None:
+        return Tensor(_infinicore.logical_or(input._underlying, other._underlying))
+
+    _infinicore.logical_or_(out._underlying, input._underlying, other._underlying)
+
+    return out

python/infinicore/ops/logical_xor.py

@@ -0,0 +1,15 @@
+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def logical_xor(input, other, *, out=None):
+    if infinicore.use_ntops and input.device.type in ("cuda", "musa"):
+        return infinicore.ntops.torch.logical_xor(input, other, out=out)
+
+    if out is None:
+        return Tensor(_infinicore.logical_xor(input._underlying, other._underlying))
+
+    _infinicore.logical_xor_(out._underlying, input._underlying, other._underlying)
+
+    return out