【比赛2025秋】T1-1-32

include/infinicore/ops/cat.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+Tensor cat(std::vector<Tensor> tensors, int dim);
+void cat_(Tensor out, std::vector<Tensor> tensors, int dim);
+} // namespace infinicore::op

include/infinicore/ops/inner.hpp

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

include/infinicore/ops/masked_select.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class MaskedSelect {
+public:
+    using schema = void (*)(Tensor, Tensor, void **, size_t *);
+    static void execute(Tensor input, Tensor mask, void **data_ptr, size_t *dlen_ptr);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor masked_select(Tensor input, Tensor mask);
+
+} // namespace infinicore::op

include/infinicore/ops/tan.hpp

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

include/infinicore/ops/tanhshrink.hpp

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

include/infiniop.h

@@ -19,20 +19,22 @@
 #include "infiniop/ops/cross_entropy.h"
 #include "infiniop/ops/dequant/per_tensor_dequant_int8.h"
 #include "infiniop/ops/dequantize_awq.h"
+#include "infiniop/ops/dequantize_gptq.h"
 #include "infiniop/ops/embedding.h"
 #include "infiniop/ops/equal.h"
 #include "infiniop/ops/flash_attention.h"
 #include "infiniop/ops/fmod.h"
-#include "infiniop/ops/dequantize_gptq.h"
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
 #include "infiniop/ops/hardswish.h"
 #include "infiniop/ops/hardtanh.h"
+#include "infiniop/ops/inner.h"
 #include "infiniop/ops/int8_gemm.h"
 #include "infiniop/ops/kv_caching.h"
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/lp_norm.h"
+#include "infiniop/ops/masked_select.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/ones.h"
 #include "infiniop/ops/paged_attention.h"
@@ -54,7 +56,9 @@
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/sum.h"
 #include "infiniop/ops/swiglu.h"
+#include "infiniop/ops/tan.h"
 #include "infiniop/ops/tanh.h"
+#include "infiniop/ops/tanhshrink.h"
 #include "infiniop/ops/topk.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"

include/infiniop/ops/inner.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_INNER_API_H__
+#define __INFINIOP_INNER_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopInnerDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateInnerDescriptor(
+    infiniopHandle_t handle,
+    infiniopInnerDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t other_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetInnerWorkspaceSize(infiniopInnerDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopInner(
+    infiniopInnerDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *out,
+    const void *input,
+    const void *other,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyInnerDescriptor(infiniopInnerDescriptor_t desc);
+
+#endif

include/infiniop/ops/masked_select.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_MASKED_SELECT_API_H__
+#define __INFINIOP_MASKED_SELECT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopMaskedSelectDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateMaskedSelectDescriptor(
+    infiniopHandle_t handle,
+    infiniopMaskedSelectDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t mask_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetMaskedSelectWorkspaceSize(infiniopMaskedSelectDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopMaskedSelect(
+    infiniopMaskedSelectDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    const void *input,
+    const bool *mask,
+    void **data_ptr,
+    size_t *dlen,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyMaskedSelectDescriptor(infiniopMaskedSelectDescriptor_t desc);
+
+#endif

include/infiniop/ops/tan.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_TAN_API_H__
+#define __INFINIOP_TAN_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTanDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateTanDescriptor(
+    infiniopHandle_t handle,
+    infiniopTanDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t out_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetTanWorkspaceSize(infiniopTanDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopTan(
+    infiniopTanDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyTanDescriptor(infiniopTanDescriptor_t desc);
+
+#endif

include/infiniop/ops/tanhshrink.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_TANHSHRINK_API_H__
+#define __INFINIOP_TANHSHRINK_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTanhshrinkDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateTanhshrinkDescriptor(infiniopHandle_t handle,
+                                                                      infiniopTanhshrinkDescriptor_t *desc_ptr,
+                                                                      infiniopTensorDescriptor_t output,
+                                                                      infiniopTensorDescriptor_t intput);
+
+__INFINI_C __export infiniStatus_t infiniopGetTanhshrinkWorkspaceSize(infiniopTanhshrinkDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopTanhshrink(infiniopTanhshrinkDescriptor_t desc,
+                                                      void *workspace,
+                                                      size_t workspace_size,
+                                                      void *output,
+                                                      const void *intput,
+                                                      void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyTanhshrinkDescriptor(infiniopTanhshrinkDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -59,11 +59,14 @@
 from infinicore.ops.binary_cross_entropy_with_logits import (
     binary_cross_entropy_with_logits,
 )
+from infinicore.ops.cat import cat
 from infinicore.ops.cdist import cdist
 from infinicore.ops.cross_entropy import cross_entropy
 from infinicore.ops.equal import equal
 from infinicore.ops.fmod import fmod
+from infinicore.ops.inner import inner
 from infinicore.ops.kv_caching import kv_caching
+from infinicore.ops.masked_select import masked_select
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mha_kvcache import mha_kvcache
 from infinicore.ops.mha_varlen import mha_varlen
@@ -76,6 +79,7 @@
 from infinicore.ops.reciprocal import reciprocal
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.sum import sum
+from infinicore.ops.tan import tan
 from infinicore.ops.topk import topk
 from infinicore.ops.unsqueeze import unsqueeze
 from infinicore.ops.var import var
@@ -151,6 +155,9 @@
     "baddbmm",
     "bilinear",
     "fmod",
+    "cat",
+    "inner",
+    "masked_select",
     "matmul",
     "equal",
     "mul",
@@ -159,6 +166,7 @@
     "unsqueeze",
     "rearrange",
     "cross_entropy",
+    "tan",
     "empty",
     "empty_like",
     "from_blob",

python/infinicore/nn/functional/__init__.py

@@ -14,6 +14,7 @@
 from .silu import silu
 from .silu_and_mul import silu_and_mul
 from .swiglu import swiglu
+from .tanhshrink import tanhshrink
 
 __all__ = [
     "adaptive_max_pool1d",
@@ -33,4 +34,5 @@
     "swiglu",
     "linear_w8a8i8",
     "silu_and_mul",
+    "tanhshrink",
 ]

python/infinicore/nn/functional/tanhshrink.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def tanhshrink(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.tanhshrink(input._underlying))
+
+    _infinicore.tanhshrink_(out._underlying, input._underlying)
+
+    return out

python/infinicore/ops/cat.py

@@ -0,0 +1,12 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def cat(tensors, dim: int = 0, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.cat([tensor._underlying for tensor in tensors], dim))
+
+    _infinicore.cat_(out._underlying, [tensor._underlying for tensor in tensors], dim)
+
+    # raise RuntimeError("breakpointer!!!")
+    return out

python/infinicore/ops/inner.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def inner(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.inner(input._underlying, other._underlying))
+
+    _infinicore.inner_(out._underlying, input._underlying, other._underlying)
+
+    return out

python/infinicore/ops/masked_select.py

@@ -0,0 +1,6 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def masked_select(input, mask):
+    return Tensor(_infinicore.masked_select(input._underlying, mask._underlying))

python/infinicore/ops/tan.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def tan(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.tan(input._underlying))
+
+    _infinicore.tan_(out._underlying, input._underlying)
+
+    return out