【比赛2025秋】T1-1-49

include/infinicore/ops/log_softmax.hpp

@@ -0,0 +1,23 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class LogSoftmax {
+public:
+    // Schema signature: output(out), input, dim
+    using schema = void (*)(Tensor, Tensor, int64_t);
+
+    static void execute(Tensor output, Tensor input, int64_t dim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API: Returns the result tensor
+Tensor log_softmax(Tensor input, int64_t dim);
+
+// In-place/Output-provided API
+void log_softmax_(Tensor output, Tensor input, int64_t dim);
+
+} // namespace infinicore::op

include/infinicore/ops/logaddexp.hpp

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

include/infinicore/ops/logaddexp2.hpp

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

include/infinicore/ops/triplet_margin_with_distance_loss.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class TripletMarginWithDistanceLoss {
+public:
+    // Schema signature: output(out), anchor, positive, negative, margin, swap, reduction
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, double, bool, int64_t);
+
+    static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API: Returns the result tensor
+// margin default 1.0, swap default false, reduction default 1 (Mean) typically
+Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin = 1.0, bool swap = false, int64_t reduction = 1);
+
+// In-place/Output-provided API
+void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction);
+
+} // namespace infinicore::op

include/infinicore/ops/upsample_nearest.hpp

@@ -0,0 +1,26 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <vector>
+
+namespace infinicore::op {
+
+class UpsampleNearest {
+public:
+    // Schema signature: output(out), input
+    // Note: Scales are inferred from output.shape / input.shape
+    using schema = void (*)(Tensor, Tensor);
+
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API: Returns the result tensor
+// Requires output_size to calculate the shape of the result tensor
+Tensor upsample_nearest(Tensor input, const std::vector<int64_t> &output_size);
+
+// In-place/Output-provided API
+void upsample_nearest_(Tensor output, Tensor input);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -50,7 +50,9 @@
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/ldexp.h"
 #include "infiniop/ops/lerp.h"
-#include "infiniop/ops/logsoftmax.h"
+#include "infiniop/ops/log_softmax.h"
+#include "infiniop/ops/logaddexp.h"
+#include "infiniop/ops/logaddexp2.h"
 #include "infiniop/ops/lp_norm.h"
 #include "infiniop/ops/masked_select.h"
 #include "infiniop/ops/mul.h"
@@ -85,7 +87,9 @@
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
 #include "infiniop/ops/triplet_margin_loss.h"
+#include "infiniop/ops/triplet_margin_with_distance_loss.h"
 #include "infiniop/ops/upsample_bilinear.h"
+#include "infiniop/ops/upsample_nearest.h"
 #include "infiniop/ops/var.h"
 #include "infiniop/ops/var_mean.h"
 #include "infiniop/ops/zeros.h"

include/infiniop/ops/logsoftmax.h → include/infiniop/ops/log_softmax.h

@@ -1,24 +1,25 @@
-#ifndef __INFINIOP_LOGSOFTMAX_API_H__
-#define __INFINIOP_LOGSOFTMAX_API_H__
+#ifndef __INFINIOP_LOG_SOFTMAX_API_H__
+#define __INFINIOP_LOG_SOFTMAX_API_H__
 
 #include "../operator_descriptor.h"
 
 typedef struct InfiniopDescriptor *infiniopLogSoftmaxDescriptor_t;
 
 __INFINI_C __export infiniStatus_t infiniopCreateLogSoftmaxDescriptor(infiniopHandle_t handle,
-                                                                  infiniopLogSoftmaxDescriptor_t *desc_ptr,
-                                                                  infiniopTensorDescriptor_t y_desc,
-                                                                  infiniopTensorDescriptor_t x_desc);
+                                                               infiniopLogSoftmaxDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t output,
+                                                               infiniopTensorDescriptor_t input,
+                                                               int dim);
 
 __INFINI_C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size);
 
 __INFINI_C __export infiniStatus_t infiniopLogSoftmax(infiniopLogSoftmaxDescriptor_t desc,
                                                void *workspace,
                                                size_t workspace_size,
-                                               void *y,
-                                               const void *x,
+                                               void *output,
+                                               const void *input,
                                                void *stream);
 
 __INFINI_C __export infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc);
 
-#endif
+#endif // __INFINIOP_LOG_SOFTMAX_API_H__

include/infiniop/ops/logaddexp.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_LOGADDEXP_API_H__
+#define __INFINIOP_LOGADDEXP_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogAddExpDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateLogAddExpDescriptor(infiniopHandle_t handle,
+                                                                     infiniopLogAddExpDescriptor_t *desc_ptr,
+                                                                     infiniopTensorDescriptor_t c,
+                                                                     infiniopTensorDescriptor_t a,
+                                                                     infiniopTensorDescriptor_t b);
+
+__INFINI_C __export infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopLogAddExp(infiniopLogAddExpDescriptor_t desc,
+                                                     void *workspace,
+                                                     size_t workspace_size,
+                                                     void *c,
+                                                     const void *a,
+                                                     const void *b,
+                                                     void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDescriptor_t desc);
+
+#endif // __INFINIOP_LOGADDEXP_API_H__

include/infiniop/ops/logaddexp2.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_LOGADDEXP2_API_H__
+#define __INFINIOP_LOGADDEXP2_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogAddExp2Descriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateLogAddExp2Descriptor(infiniopHandle_t handle,
+                                                                      infiniopLogAddExp2Descriptor_t *desc_ptr,
+                                                                      infiniopTensorDescriptor_t c,
+                                                                      infiniopTensorDescriptor_t a,
+                                                                      infiniopTensorDescriptor_t b);
+
+__INFINI_C __export infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopLogAddExp2(infiniopLogAddExp2Descriptor_t desc,
+                                                      void *workspace,
+                                                      size_t workspace_size,
+                                                      void *c,
+                                                      const void *a,
+                                                      const void *b,
+                                                      void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2Descriptor_t desc);
+
+#endif // __INFINIOP_LOGADDEXP2_API_H__

include/infiniop/ops/triplet_margin_with_distance_loss.h

@@ -0,0 +1,32 @@
+#ifndef __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__
+#define __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTripletMarginWithDistanceLossDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor(
+    infiniopHandle_t handle,
+    infiniopTripletMarginWithDistanceLossDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output,
+    infiniopTensorDescriptor_t anchor,
+    infiniopTensorDescriptor_t positive,
+    infiniopTensorDescriptor_t negative,
+    float margin,
+    int swap,
+    int reduction);
+__INFINI_C __export infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize(
+    infiniopTripletMarginWithDistanceLossDescriptor_t desc,
+    size_t *size);
+__INFINI_C __export infiniStatus_t infiniopTripletMarginWithDistanceLoss(infiniopTripletMarginWithDistanceLossDescriptor_t desc,
+                                                                         void *workspace,
+                                                                         size_t workspace_size,
+                                                                         void *output,
+                                                                         const void *anchor,
+                                                                         const void *positive,
+                                                                         const void *negative,
+                                                                         void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyTripletMarginWithDistanceLossDescriptor(
+    infiniopTripletMarginWithDistanceLossDescriptor_t desc);
+#endif // __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__

include/infiniop/ops/upsample_nearest.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_UPSAMPLE_NEAREST_API_H__
+#define __INFINIOP_UPSAMPLE_NEAREST_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopUpsampleNearestDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateUpsampleNearestDescriptor(infiniopHandle_t handle,
+                                                                           infiniopUpsampleNearestDescriptor_t *desc_ptr,
+                                                                           infiniopTensorDescriptor_t output,
+                                                                           infiniopTensorDescriptor_t input);
+
+__INFINI_C __export infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopUpsampleNearest(infiniopUpsampleNearestDescriptor_t desc,
+                                                           void *workspace,
+                                                           size_t workspace_size,
+                                                           void *output,
+                                                           const void *input,
+                                                           void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc);
+
+#endif // __INFINIOP_UPSAMPLE_NEAREST_API_H__

python/infinicore/__init__.py

@@ -82,6 +82,8 @@
 from infinicore.ops.kv_caching import kv_caching
 from infinicore.ops.ldexp import ldexp
 from infinicore.ops.lerp import lerp
+from infinicore.ops.logaddexp import logaddexp
+from infinicore.ops.logaddexp2 import logaddexp2
 from infinicore.ops.masked_select import masked_select
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mha_kvcache import mha_kvcache
@@ -182,6 +184,8 @@
     "cat",
     "inner",
     "masked_select",
+    "logaddexp",
+    "logaddexp2",
     "matmul",
     "equal",
     "mul",

python/infinicore/nn/functional/__init__.py

@@ -9,8 +9,10 @@
 from .flash_attention import flash_attention
 from .hardswish import hardswish
 from .hardtanh import hardtanh
+from .interpolate import interpolate
 from .linear import linear
 from .linear_w8a8i8 import linear_w8a8i8
+from .log_softmax import log_softmax
 from .multi_margin_loss import multi_margin_loss
 from .random_sample import random_sample
 from .rms_norm import rms_norm
@@ -21,7 +23,9 @@
 from .swiglu import swiglu
 from .tanhshrink import tanhshrink
 from .triplet_margin_loss import triplet_margin_loss
-from .upsample_bilinear import interpolate, upsample_bilinear
+from .triplet_margin_with_distance_loss import triplet_margin_with_distance_loss
+from .upsample_bilinear import upsample_bilinear
+from .upsample_nearest import upsample_nearest
 
 __all__ = [
     "adaptive_max_pool1d",
@@ -37,10 +41,14 @@
     "silu",
     "smooth_l1_loss",
     "swiglu",
+    "interpolate",
     "linear",
     "triplet_margin_loss",
     "upsample_bilinear",
     "interpolate",
+    "log_softmax",
+    "upsample_nearest",
+    "triplet_margin_with_distance_loss",
     "embedding",
     "rope",
     "RopeAlgo",

python/infinicore/nn/functional/interpolate.py

@@ -0,0 +1,32 @@
+from typing import Optional, Sequence, Union
+
+from infinicore.tensor import Tensor
+
+from .upsample_bilinear import upsample_bilinear
+from .upsample_nearest import upsample_nearest
+
+
+def interpolate(
+    input: Tensor,
+    size: Optional[Union[int, Sequence[int]]] = None,
+    scale_factor: Optional[Union[float, Sequence[float]]] = None,
+    mode: str = "nearest",
+    align_corners: Optional[bool] = None,
+    recompute_scale_factor: Optional[bool] = None,
+) -> Tensor:
+    if mode == "nearest":
+        if align_corners is not None:
+            raise ValueError(
+                "align_corners option can only be set with the "
+                "interpolating modes: linear | bilinear | bicubic | trilinear"
+            )
+        return upsample_nearest(input, size, scale_factor)
+
+    if mode == "bilinear":
+        if align_corners is None:
+            align_corners = False
+        return upsample_bilinear(input, size, scale_factor, align_corners)
+
+    raise NotImplementedError(
+        f"Interpolation mode '{mode}' is not currently supported."
+    )

python/infinicore/nn/functional/log_softmax.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def log_softmax(input: Tensor, dim: int, *, out: Optional[Tensor] = None) -> Tensor:
+    r"""Applies a softmax followed by a logarithm.
+    While mathematically equivalent to log(softmax(x)), doing these two
+    operations separately is slower and numerically unstable. This function
+    uses an alternative formulation to compute the output and gradient correctly.
+    """
+
+    if not input.is_contiguous():
+        input = input.contiguous()
+
+    if out is not None:
+        if not isinstance(out, Tensor):
+            raise ValueError("out must be a Tensor")
+
+        _infinicore.log_softmax_(out._underlying, input._underlying, dim)
+        return out
+
+    ret = _infinicore.log_softmax(input._underlying, dim)
+
+    return Tensor(ret)