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

include/infinicore/ops/kthvalue.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <tuple>
+
+namespace infinicore::op {
+
+class Kthvalue {
+public:
+    // Schema signature: values(out), indices(out), input, k, dim, keepdim
+    using schema = void (*)(Tensor, Tensor, Tensor, int64_t, int64_t, bool);
+
+    static void execute(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API: Returns a tuple containing (values, indices)
+std::tuple<Tensor, Tensor> kthvalue(Tensor input, int64_t k, int64_t dim = -1, bool keepdim = false);
+
+// In-place/Output-provided API
+void kthvalue_(Tensor values, Tensor indices, Tensor input, int64_t k, int64_t dim, bool keepdim);
+
+} // namespace infinicore::op

include/infinicore/ops/ldexp.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Ldexp {
+public:
+    // Schema signature: output(out), input(x), other(exp)
+    using schema = void (*)(Tensor, Tensor, Tensor);
+
+    static void execute(Tensor output, Tensor input, Tensor other);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API: Returns a new Tensor containing input * (2^other)
+Tensor ldexp(Tensor input, Tensor other);
+
+// In-place/Output-provided API
+// Writes the result into 'output'
+void ldexp_(Tensor output, Tensor input, Tensor other);
+
+} // namespace infinicore::op

include/infinicore/ops/lerp.hpp

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Lerp {
+public:
+    using schema_t = void (*)(Tensor, Tensor, Tensor, Tensor);
+    using schema_s = void (*)(Tensor, Tensor, Tensor, float);
+
+    static void execute(Tensor output, Tensor start, Tensor end, Tensor weight);
+    static void execute(Tensor output, Tensor start, Tensor end, float weight);
+
+    // 【核心修改】必须声明为模板函数，才能支持 dispatcher<schema_t>() 和 dispatcher<schema_s>()
+    template <typename T>
+    static common::OpDispatcher<T> &dispatcher();
+};
+
+Tensor lerp(Tensor start, Tensor end, Tensor weight);
+Tensor lerp(Tensor start, Tensor end, float weight);
+
+void lerp_(Tensor output, Tensor start, Tensor end, Tensor weight);
+void lerp_(Tensor output, Tensor start, Tensor end, float weight);
+
+} // namespace infinicore::op

include/infinicore/ops/triplet_margin_loss.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class TripletMarginLoss {
+public:
+    // Schema signature: output, anchor, positive, negative, margin, p, eps, swap, reduction
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, int64_t, float, bool, int64_t);
+
+    static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// Functional API
+// reduction: 0=None, 1=Mean, 2=Sum
+Tensor triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin = 1.0f, int64_t p = 2, float eps = 1e-6f, bool swap = false, int64_t reduction = 1);
+
+// In-place / Explicit Output API
+void triplet_margin_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, float margin, int64_t p, float eps, bool swap, int64_t reduction);
+
+} // namespace infinicore::op

include/infinicore/ops/upsample_bilinear.hpp

@@ -0,0 +1,22 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <vector>
+
+namespace infinicore::op {
+
+class UpsampleBilinear {
+public:
+    // Schema signature: output, input, align_corners
+    using schema = void (*)(Tensor, Tensor, bool);
+
+    static void execute(Tensor output, Tensor input, bool align_corners);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// 需要传入 output_size (如 {H_out, W_out} 或 {N, C, H_out, W_out}) 来决定新 Tensor 的形状
+Tensor upsample_bilinear(Tensor input, std::vector<int64_t> output_size, bool align_corners = false);
+void upsample_bilinear_(Tensor output, Tensor input, bool align_corners);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -29,7 +29,11 @@
 #include "infiniop/ops/tanh.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
+#include "infiniop/ops/triplet_margin_loss.h"
+#include "infiniop/ops/upsample_bilinear.h"
+#include "infiniop/ops/kthvalue.h"
+#include "infiniop/ops/lerp.h"
+#include "infiniop/ops/ldexp.h"
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
-
 #endif // __INFINIOP_API_H__

include/infiniop/ops/kthvalue.h

@@ -0,0 +1,29 @@
+#ifndef __INFINIOP_KTHVALUE_API_H__
+#define __INFINIOP_KTHVALUE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopKthvalueDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateKthvalueDescriptor(infiniopHandle_t handle,
+                                                             infiniopKthvalueDescriptor_t *desc_ptr,
+                                                             infiniopTensorDescriptor_t values,
+                                                             infiniopTensorDescriptor_t indices,
+                                                             infiniopTensorDescriptor_t input,
+                                                             int k,
+                                                             int dim,
+                                                             int keepdim);
+
+__C __export infiniStatus_t infiniopGetKthvalueWorkspaceSize(infiniopKthvalueDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopKthvalue(infiniopKthvalueDescriptor_t desc,
+                                             void *workspace,
+                                             size_t workspace_size,
+                                             void *values,
+                                             void *indices,
+                                             const void *input,
+                                             void *stream);
+
+__C __export infiniStatus_t infiniopDestroyKthvalueDescriptor(infiniopKthvalueDescriptor_t desc);
+
+#endif // __INFINIOP_KTHVALUE_API_H__

include/infiniop/ops/ldexp.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_LDEXP_API_H__
+#define __INFINIOP_LDEXP_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLdexpDescriptor_t;
+__C __export infiniStatus_t infiniopCreateLdexpDescriptor(infiniopHandle_t handle,
+                                                          infiniopLdexpDescriptor_t *desc_ptr,
+                                                          infiniopTensorDescriptor_t y,
+                                                          infiniopTensorDescriptor_t x,
+                                                          infiniopTensorDescriptor_t exp);
+__C __export infiniStatus_t infiniopGetLdexpWorkspaceSize(infiniopLdexpDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLdexp(infiniopLdexpDescriptor_t desc,
+                                          void *workspace,
+                                          size_t workspace_size,
+                                          void *y,
+                                          const void *x,
+                                          const void *exp,
+                                          void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLdexpDescriptor(infiniopLdexpDescriptor_t desc);
+
+#endif // __INFINIOP_LDEXP_API_H__

include/infiniop/ops/lerp.h

@@ -0,0 +1,29 @@
+#ifndef __INFINIOP_LERP_API_H__
+#define __INFINIOP_LERP_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLerpDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLerpDescriptor(infiniopHandle_t handle,
+                                                         infiniopLerpDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t start,
+                                                         infiniopTensorDescriptor_t end,
+                                                         infiniopTensorDescriptor_t weight,
+                                                         float weight_scalar);
+
+__C __export infiniStatus_t infiniopGetLerpWorkspaceSize(infiniopLerpDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLerp(infiniopLerpDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *start,
+                                         const void *end,
+                                         const void *weight,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLerpDescriptor(infiniopLerpDescriptor_t desc);
+
+#endif // __INFINIOP_LERP_API_H__

include/infiniop/ops/triplet_margin_loss.h

@@ -0,0 +1,33 @@
+#ifndef __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__
+#define __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTripletMarginLossDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTripletMarginLossDescriptor(infiniopHandle_t handle,
+                                                                      infiniopTripletMarginLossDescriptor_t *desc_ptr,
+                                                                      infiniopTensorDescriptor_t output,
+                                                                      infiniopTensorDescriptor_t anchor,
+                                                                      infiniopTensorDescriptor_t positive,
+                                                                      infiniopTensorDescriptor_t negative,
+                                                                      float margin,
+                                                                      int p,
+                                                                      float eps,
+                                                                      int swap,      // 0: False, 1: True
+                                                                      int reduction); // 0: None, 1: Mean, 2: Sum
+
+__C __export infiniStatus_t infiniopGetTripletMarginLossWorkspaceSize(infiniopTripletMarginLossDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTripletMarginLoss(infiniopTripletMarginLossDescriptor_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 infiniopDestroyTripletMarginLossDescriptor(infiniopTripletMarginLossDescriptor_t desc);
+
+#endif // __INFINIOP_TRIPLET_MARGIN_LOSS_API_H__

include/infiniop/ops/upsample_bilinear.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_UPSAMPLE_BILINEAR_API_H__
+#define __INFINIOP_UPSAMPLE_BILINEAR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopUpsampleBilinearDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateUpsampleBilinearDescriptor(infiniopHandle_t handle,
+                                                                     infiniopUpsampleBilinearDescriptor_t *desc_ptr,
+                                                                     infiniopTensorDescriptor_t output,
+                                                                     infiniopTensorDescriptor_t input,
+                                                                     int align_corners);
+
+__C __export infiniStatus_t infiniopGetUpsampleBilinearWorkspaceSize(infiniopUpsampleBilinearDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopUpsampleBilinear(infiniopUpsampleBilinearDescriptor_t desc,
+                                                     void *workspace,
+                                                     size_t workspace_size,
+                                                     void *output,
+                                                     const void *input,
+                                                     void *stream);
+
+__C __export infiniStatus_t infiniopDestroyUpsampleBilinearDescriptor(infiniopUpsampleBilinearDescriptor_t desc);
+
+#endif // __INFINIOP_UPSAMPLE_BILINEAR_API_H__

python/infinicore/__init__.py

@@ -44,6 +44,9 @@
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
+from infinicore.ops.ldexp import ldexp
+from infinicore.ops.lerp import lerp
+from infinicore.ops.kthvalue import kthvalue
 from infinicore.ops.rearrange import rearrange
 from infinicore.tensor import (
     Tensor,
@@ -104,6 +107,9 @@
     "matmul",
     "mul",
     "narrow",
+    "ldexp",    
+    "lerp",      
+    "kthvalue",  
     "rearrange",
     "empty",
     "empty_like",

python/infinicore/nn/functional/__init__.py

@@ -6,14 +6,18 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
-
+from .upsample_bilinear import upsample_bilinear, interpolate
+from .triplet_margin_loss import triplet_margin_loss 
 __all__ = [
     "causal_softmax",
     "random_sample",
     "rms_norm",
     "silu",
     "swiglu",
     "linear",
+    "triplet_margin_loss",
+    "upsample_bilinear",
+    "interpolate", 
     "embedding",
     "rope",
     "RopeAlgo",

python/infinicore/nn/functional/triplet_margin_loss.py

@@ -0,0 +1,63 @@
+from typing import Optional
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+_REDUCTION_MODES = {
+    "none": 0,
+    "mean": 1,
+    "sum": 2,
+}
+
+def triplet_margin_loss(
+    anchor: Tensor, 
+    positive: Tensor, 
+    negative: Tensor, 
+    margin: float = 1.0, 
+    p: float = 2, 
+    eps: float = 1e-6, 
+    swap: bool = False, 
+    reduction: str = "mean", 
+    *, 
+    out: Optional[Tensor] = None
+) -> Tensor:
+    r"""Creates a criterion that measures the triplet loss given an input
+    tensors x1, x2, x3 and a margin with a value greater than 0.
+    """
+
+    if not anchor.is_contiguous():
+        anchor = anchor.contiguous()
+    if not positive.is_contiguous():
+        positive = positive.contiguous()
+    if not negative.is_contiguous():
+        negative = negative.contiguous()
+
+    if reduction not in _REDUCTION_MODES:
+        raise ValueError(f"{reduction} is not a valid value for reduction")
+    reduction_val = _REDUCTION_MODES[reduction]
+
+    if out is not None:
+        _infinicore.triplet_margin_loss_(
+            out._underlying,
+            anchor._underlying,
+            positive._underlying,
+            negative._underlying,
+            margin,
+            int(p),
+            eps,
+            swap,
+            reduction_val
+        )
+        return out
+
+    return Tensor(
+        _infinicore.triplet_margin_loss(
+            anchor._underlying,
+            positive._underlying,
+            negative._underlying,
+            margin,
+            int(p),
+            eps,
+            swap,
+            reduction_val
+        )
+    )