【比赛2025秋】T1-1-30

include/infinicore/ops/bitwise_right_shift.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(BitwiseRightShift, Tensor, const Tensor &, const Tensor &);
+
+__export Tensor bitwise_right_shift(const Tensor &input, const Tensor &other);
+__export void bitwise_right_shift_(Tensor out, const Tensor &input, const Tensor &other);
+
+} // namespace infinicore::op

include/infinicore/ops/gaussian_nll_loss.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(GaussianNllLoss, Tensor, const Tensor &, const Tensor &, const Tensor &, bool, double, int);
+
+__export Tensor gaussian_nll_loss(const Tensor &input,
+                                  const Tensor &target,
+                                  const Tensor &var,
+                                  bool full = false,
+                                  double eps = 1e-6,
+                                  int reduction = 1);
+
+__export void gaussian_nll_loss_(Tensor out,
+                                 const Tensor &input,
+                                 const Tensor &target,
+                                 const Tensor &var,
+                                 bool full = false,
+                                 double eps = 1e-6,
+                                 int reduction = 1);
+
+} // namespace infinicore::op

include/infinicore/ops/interpolate.hpp

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Interpolate, Tensor, const Tensor &, std::string, std::vector<int64_t>, std::vector<double>, int);
+
+__export Tensor interpolate(const Tensor &input,
+                            std::string mode,
+                            std::vector<int64_t> size,
+                            std::vector<double> scale_factor,
+                            int align_corners);
+
+__export void interpolate_(Tensor out,
+                           const Tensor &input,
+                           std::string mode,
+                           std::vector<int64_t> size,
+                           std::vector<double> scale_factor,
+                           int align_corners);
+
+} // namespace infinicore::op

include/infinicore/ops/prelu.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Prelu, Tensor, const Tensor &, const Tensor &);
+
+__export Tensor prelu(const Tensor &input, const Tensor &weight);
+__export void prelu_(Tensor out, const Tensor &input, const Tensor &weight);
+
+} // namespace infinicore::op

include/infinicore/ops/relu6.hpp

@@ -0,0 +1,16 @@
+#pragma once
+
+#include "infinicore.h"
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(Relu6, Tensor, const Tensor &);
+
+__export Tensor relu6(const Tensor &input);
+__export void relu6_(Tensor out, const Tensor &input);
+
+} // namespace infinicore::op

include/infiniop/ops/bitwise_right_shift.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__
+#define __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopBitwiseRightShiftDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor(infiniopHandle_t handle,
+                                                                             infiniopBitwiseRightShiftDescriptor_t *desc_ptr,
+                                                                             infiniopTensorDescriptor_t y,
+                                                                             infiniopTensorDescriptor_t x1,
+                                                                             infiniopTensorDescriptor_t x2);
+
+__INFINI_C __export infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitwiseRightShiftDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopBitwiseRightShift(infiniopBitwiseRightShiftDescriptor_t desc,
+                                                             void *workspace,
+                                                             size_t workspace_size,
+                                                             void *y,
+                                                             const void *x1,
+                                                             const void *x2,
+                                                             void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t desc);
+
+#endif

include/infiniop/ops/gaussian_nll_loss.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__
+#define __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopGaussianNllLossDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateGaussianNllLossDescriptor(infiniopHandle_t handle,
+                                                                           infiniopGaussianNllLossDescriptor_t *desc_ptr,
+                                                                           infiniopTensorDescriptor_t y,
+                                                                           infiniopTensorDescriptor_t input,
+                                                                           infiniopTensorDescriptor_t target,
+                                                                           infiniopTensorDescriptor_t var,
+                                                                           int full,
+                                                                           double eps,
+                                                                           int reduction);
+
+__INFINI_C __export infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussianNllLossDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopGaussianNllLoss(infiniopGaussianNllLossDescriptor_t desc,
+                                                           void *workspace,
+                                                           size_t workspace_size,
+                                                           void *y,
+                                                           const void *input,
+                                                           const void *target,
+                                                           const void *var,
+                                                           void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t desc);
+
+#endif

include/infiniop/ops/interpolate.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_INTERPOLATE_API_H__
+#define __INFINIOP_INTERPOLATE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopInterpolateDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateInterpolateDescriptor(infiniopHandle_t handle,
+                                                                       infiniopInterpolateDescriptor_t *desc_ptr,
+                                                                       infiniopTensorDescriptor_t y,
+                                                                       infiniopTensorDescriptor_t x,
+                                                                       const char *mode,
+                                                                       void *size,
+                                                                       void *scale_factor,
+                                                                       int align_corners);
+
+__INFINI_C __export infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolateDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopInterpolate(infiniopInterpolateDescriptor_t desc,
+                                                       void *workspace,
+                                                       size_t workspace_size,
+                                                       void *y,
+                                                       const void *x,
+                                                       void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc);
+
+#endif

include/infiniop/ops/prelu.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_PRELU_API_H__
+#define __INFINIOP_PRELU_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopPreluDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreatePreluDescriptor(infiniopHandle_t handle,
+                                                                 infiniopPreluDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t y,
+                                                                 infiniopTensorDescriptor_t x,
+                                                                 infiniopTensorDescriptor_t weight);
+
+__INFINI_C __export infiniStatus_t infiniopGetPreluWorkspaceSize(infiniopPreluDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopPrelu(infiniopPreluDescriptor_t desc,
+                                                 void *workspace,
+                                                 size_t workspace_size,
+                                                 void *y,
+                                                 const void *x,
+                                                 const void *weight,
+                                                 void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyPreluDescriptor(infiniopPreluDescriptor_t desc);
+
+#endif

include/infiniop/ops/relu6.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_RELU6_API_H__
+#define __INFINIOP_RELU6_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopRelu6Descriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateRelu6Descriptor(infiniopHandle_t handle,
+                                                                 infiniopRelu6Descriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t y,
+                                                                 infiniopTensorDescriptor_t x);
+
+__INFINI_C __export infiniStatus_t infiniopGetRelu6WorkspaceSize(infiniopRelu6Descriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopRelu6(infiniopRelu6Descriptor_t desc,
+                                                 void *workspace,
+                                                 size_t workspace_size,
+                                                 void *y,
+                                                 const void *x,
+                                                 void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyRelu6Descriptor(infiniopRelu6Descriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -64,6 +64,7 @@
 from infinicore.ops.binary_cross_entropy_with_logits import (
     binary_cross_entropy_with_logits,
 )
+from infinicore.ops.bitwise_right_shift import bitwise_right_shift
 from infinicore.ops.block_diag import block_diag
 from infinicore.ops.broadcast_to import broadcast_to
 from infinicore.ops.cat import cat
@@ -189,6 +190,7 @@
     "attention",
     "block_diag",
     "kron",
+    "bitwise_right_shift",
     "kv_caching",
     "asinh",
     "baddbmm",

python/infinicore/nn/functional/__init__.py

@@ -7,6 +7,7 @@
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .flash_attention import flash_attention
+from .gaussian_nll_loss import gaussian_nll_loss
 from .hardswish import hardswish
 from .hardtanh import hardtanh
 from .hinge_embedding_loss import hinge_embedding_loss
@@ -17,7 +18,9 @@
 from .log_softmax import log_softmax
 from .multi_margin_loss import multi_margin_loss
 from .pad import pad
+from .prelu import prelu
 from .random_sample import random_sample
+from .relu6 import relu6
 from .rms_norm import rms_norm
 from .rope import RopeAlgo, rope
 from .selu import selu
@@ -38,11 +41,15 @@
     "causal_softmax",
     "embedding",
     "flash_attention",
+    "gaussian_nll_loss",
+    "interpolate",
     "linear",
     "binary_cross_entropy_with_logits",
     "random_sample",
     "adaptive_avg_pool1d",
     "affine_grid",
+    "prelu",
+    "relu6",
     "rms_norm",
     "silu",
     "smooth_l1_loss",

python/infinicore/nn/functional/gaussian_nll_loss.py

@@ -0,0 +1,32 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+_REDUCTION_TO_INT = {
+    "none": 0,
+    "mean": 1,
+    "sum": 2,
+}
+
+
+def gaussian_nll_loss(
+    input: Tensor,
+    target: Tensor,
+    var: Tensor,
+    full: bool = False,
+    eps: float = 1e-6,
+    reduction: str = "mean",
+) -> Tensor:
+    reduction_i = _REDUCTION_TO_INT.get(reduction)
+    if reduction_i is None:
+        raise ValueError(f"Unsupported reduction: {reduction!r}")
+
+    return Tensor(
+        _infinicore.gaussian_nll_loss(
+            input._underlying,
+            target._underlying,
+            var._underlying,
+            bool(full),
+            float(eps),
+            int(reduction_i),
+        )
+    )