kthvalue triplet_margin_loss upsample_bilinear ldexp lerp operator

Author: lengmuzhaxi

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__