[2025秋季][T1-1-49]zili2004

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

@@ -12,8 +12,12 @@
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
 #include "infiniop/ops/layer_norm.h"
-#include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/lp_norm.h"
+#include "infiniop/ops/log_softmax.h"
+#include "infiniop/ops/upsample_nearest.h" 
+#include "infiniop/ops/triplet_margin_with_distance_loss.h"
+#include "infiniop/ops/logaddexp.h"
+#include "infiniop/ops/logaddexp2.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/ones.h"
 #include "infiniop/ops/paged_attention.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;
 
 __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);
 
 __C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size);
 
 __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);
 
 __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;
+
+__C __export infiniStatus_t infiniopCreateLogAddExpDescriptor(infiniopHandle_t handle,
+                                                              infiniopLogAddExpDescriptor_t *desc_ptr,
+                                                              infiniopTensorDescriptor_t c,
+                                                              infiniopTensorDescriptor_t a,
+                                                              infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogAddExp(infiniopLogAddExpDescriptor_t desc,
+                                              void *workspace,
+                                              size_t workspace_size,
+                                              void *c,
+                                              const void *a,
+                                              const void *b,
+                                              void *stream);
+
+__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;
+
+__C __export infiniStatus_t infiniopCreateLogAddExp2Descriptor(infiniopHandle_t handle,
+                                                               infiniopLogAddExp2Descriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t c,
+                                                               infiniopTensorDescriptor_t a,
+                                                               infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogAddExp2(infiniopLogAddExp2Descriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *c,
+                                               const void *a,
+                                               const void *b,
+                                               void *stream);
+
+__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;
+
+__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);
+__C __export infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize(
+                                                                infiniopTripletMarginWithDistanceLossDescriptor_t desc, 
+                                                                size_t *size);
+__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);
+
+__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;
+
+__C __export infiniStatus_t infiniopCreateUpsampleNearestDescriptor(infiniopHandle_t handle,
+                                                                    infiniopUpsampleNearestDescriptor_t *desc_ptr,
+                                                                    infiniopTensorDescriptor_t output,
+                                                                    infiniopTensorDescriptor_t input);
+
+__C __export infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopUpsampleNearest(infiniopUpsampleNearestDescriptor_t desc,
+                                                    void *workspace,
+                                                    size_t workspace_size,
+                                                    void *output,
+                                                    const void *input,
+                                                    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc);
+
+#endif // __INFINIOP_UPSAMPLE_NEAREST_API_H__

python/infinicore/__init__.py

@@ -42,6 +42,8 @@
 from infinicore.ops.add import add
 from infinicore.ops.add_rms_norm import add_rms_norm, add_rms_norm_
 from infinicore.ops.attention import attention
+from infinicore.ops.logaddexp2 import logaddexp2
+from infinicore.ops.logaddexp import logaddexp
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
@@ -109,6 +111,8 @@
     "add_rms_norm",
     "add_rms_norm_",
     "attention",
+    "logaddexp",
+    "logaddexp2",
     "matmul",
     "mul",
     "narrow",

python/infinicore/nn/functional/__init__.py

@@ -1,19 +1,25 @@
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .linear import linear
+from .log_softmax import log_softmax
 from .random_sample import random_sample
 from .rms_norm import rms_norm
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
-
+from .triplet_margin_with_distance_loss import triplet_margin_with_distance_loss
+from .upsample_nearest import upsample_nearest, interpolate
 __all__ = [
     "causal_softmax",
     "random_sample",
     "rms_norm",
     "silu",
     "swiglu",
+    "interpolate", 
     "linear",
+    "log_softmax",
+    "upsample_nearest",
+    "triplet_margin_with_distance_loss",
     "embedding",
     "rope",
     "RopeAlgo",

python/infinicore/nn/functional/log_softmax.py

@@ -0,0 +1,36 @@
+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)