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

include/infinicore/ops/hypot.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+class Hypot {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+
+    static void execute(Tensor output, Tensor input_a, Tensor input_b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor hypot(Tensor input_a, Tensor input_b);
+
+void hypot_(Tensor output, Tensor input_a, Tensor input_b);
+} // namespace infinicore::op

include/infinicore/ops/index_add.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class IndexAdd {
+public:
+    using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor, float);
+    static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha);
+
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+
+Tensor index_add(Tensor input, int64_t dim, Tensor index, Tensor source, float alpha = 1.0f);
+void index_add_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source, float alpha);
+
+} // namespace infinicore::op

include/infinicore/ops/index_copy.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class IndexCopy {
+public:
+    using schema = void (*)(Tensor, Tensor, int64_t, Tensor, Tensor);
+    static void execute(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source);
+
+    static common::OpDispatcher<schema> &dispatcher();
+};
+Tensor index_copy(Tensor input, int64_t dim, Tensor index, Tensor source);
+void index_copy_(Tensor output, Tensor input, int64_t dim, Tensor index, Tensor source);
+
+} // namespace infinicore::op

include/infinicore/ops/smooth_l1_loss.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class SmoothL1Loss {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, float, int64_t);
+
+    static void execute(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor smooth_l1_loss(Tensor input, Tensor target, float beta = 1.0f, int64_t reduction = 1);
+void smooth_l1_loss_(Tensor output, Tensor input, Tensor target, float beta, int64_t reduction);
+
+} // namespace infinicore::op

include/infinicore/ops/take.hpp

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Take {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+
+    static void execute(Tensor output, Tensor input, Tensor indices);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+
+Tensor take(Tensor input, Tensor indices);
+
+void take_(Tensor output, Tensor input, Tensor indices);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -10,6 +10,9 @@
 #include "infiniop/ops/dequantize_awq.h"
 #include "infiniop/ops/gelu.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/hypot.h"
+#include "infiniop/ops/index_add.h"
+#include "infiniop/ops/index_copy.h"
 #include "infiniop/ops/layer_norm.h"
 #include "infiniop/ops/logsoftmax.h"
 #include "infiniop/ops/lp_norm.h"
@@ -22,11 +25,13 @@
 #include "infiniop/ops/rope.h"
 #include "infiniop/ops/sigmoid.h"
 #include "infiniop/ops/silu.h"
+#include "infiniop/ops/smooth_l1_loss.h"
 #include "infiniop/ops/softmax.h"
 #include "infiniop/ops/softplus.h"
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/ops/tanh.h"
+#include "infiniop/ops/take.h"
 #include "infiniop/ops/topkrouter.h"
 #include "infiniop/ops/topksoftmax.h"
 #include "infiniop/ops/zeros.h"

include/infiniop/ops/hypot.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_HYPOT_API_H__
+#define __INFINIOP_HYPOT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopHypotDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateHypotDescriptor(infiniopHandle_t handle,
+                                                          infiniopHypotDescriptor_t *desc_ptr,
+                                                          infiniopTensorDescriptor_t output,
+                                                          infiniopTensorDescriptor_t input_a,
+                                                          infiniopTensorDescriptor_t input_b);
+
+__C __export infiniStatus_t infiniopGetHypotWorkspaceSize(infiniopHypotDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopHypot(infiniopHypotDescriptor_t desc,
+                                          void *workspace,
+                                          size_t workspace_size,
+                                          void *output,
+                                          const void *input_a,
+                                          const void *input_b,
+                                          void *stream);
+
+__C __export infiniStatus_t infiniopDestroyHypotDescriptor(infiniopHypotDescriptor_t desc);
+
+#endif

include/infiniop/ops/index_add.h

@@ -0,0 +1,29 @@
+#ifndef __INFINIOP_INDEX_ADD_API_H__
+#define __INFINIOP_INDEX_ADD_API_H__
+#include <stdint.h>
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopIndexAddDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateIndexAddDescriptor(infiniopHandle_t handle,
+                                                             infiniopIndexAddDescriptor_t *desc_ptr,
+                                                             infiniopTensorDescriptor_t output,
+                                                             infiniopTensorDescriptor_t input,
+                                                             int64_t dim,
+                                                             infiniopTensorDescriptor_t index,
+                                                             infiniopTensorDescriptor_t source,
+                                                             float alpha);
+
+__C __export infiniStatus_t infiniopGetIndexAddWorkspaceSize(infiniopIndexAddDescriptor_t desc, size_t *size);
+__C __export infiniStatus_t infiniopIndexAdd(infiniopIndexAddDescriptor_t desc,
+                                             void *workspace,
+                                             size_t workspace_size,
+                                             void *output,
+                                             const void *input,
+                                             const void *index,
+                                             const void *source,
+                                             void *stream);
+
+__C __export infiniStatus_t infiniopDestroyIndexAddDescriptor(infiniopIndexAddDescriptor_t desc);
+
+#endif

include/infiniop/ops/index_copy.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_INDEX_COPY_API_H__
+#define __INFINIOP_INDEX_COPY_API_H__
+#include <stdint.h>
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopIndexCopyDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateIndexCopyDescriptor(infiniopHandle_t handle,
+                                                             infiniopIndexCopyDescriptor_t *desc_ptr,
+                                                             infiniopTensorDescriptor_t output,
+                                                             infiniopTensorDescriptor_t input,
+                                                             int64_t dim,
+                                                             infiniopTensorDescriptor_t index,
+                                                             infiniopTensorDescriptor_t source);
+
+__C __export infiniStatus_t infiniopGetIndexCopyWorkspaceSize(infiniopIndexCopyDescriptor_t desc, size_t *size);
+__C __export infiniStatus_t infiniopIndexCopy(infiniopIndexCopyDescriptor_t desc,
+                                             void *workspace,
+                                             size_t workspace_size,
+                                             void *output,
+                                             const void *input,
+                                             const void *index,
+                                             const void *source,
+                                             void *stream);
+
+__C __export infiniStatus_t infiniopDestroyIndexCopyDescriptor(infiniopIndexCopyDescriptor_t desc);
+
+#endif

include/infiniop/ops/smooth_l1_loss.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_SMOOTH_L1_LOSS_API_H__
+#define __INFINIOP_SMOOTH_L1_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopSmoothL1LossDescriptor_t;
+__C __export infiniStatus_t infiniopCreateSmoothL1LossDescriptor(infiniopHandle_t handle,
+                                                                 infiniopSmoothL1LossDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t output,
+                                                                 infiniopTensorDescriptor_t input,
+                                                                 infiniopTensorDescriptor_t target,
+                                                                 float beta,
+                                                                 int reduction);
+
+__C __export infiniStatus_t infiniopGetSmoothL1LossWorkspaceSize(infiniopSmoothL1LossDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopSmoothL1Loss(infiniopSmoothL1LossDescriptor_t desc,
+                                                 void *workspace,
+                                                 size_t workspace_size,
+                                                 void *output,
+                                                 const void *input,
+                                                 const void *target,
+                                                 void *stream);
+
+__C __export infiniStatus_t infiniopDestroySmoothL1LossDescriptor(infiniopSmoothL1LossDescriptor_t desc);
+
+#endif // __INFINIOP_SMOOTH_L1_LOSS_API_H__

include/infiniop/ops/take.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_TAKE_API_H__
+#define __INFINIOP_TAKE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTakeDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTakeDescriptor(infiniopHandle_t handle,
+                                                         infiniopTakeDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t input,
+                                                         infiniopTensorDescriptor_t indices);
+
+__C __export infiniStatus_t infiniopGetTakeWorkspaceSize(infiniopTakeDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTake(infiniopTakeDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *input,
+                                         const void *indices,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyTakeDescriptor(infiniopTakeDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -45,6 +45,10 @@
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
 from infinicore.ops.rearrange import rearrange
+from infinicore.ops.hypot import hypot
+from infinicore.ops.index_add import index_add
+from infinicore.ops.index_copy import index_copy
+from infinicore.ops.take import take
 from infinicore.tensor import (
     Tensor,
     empty,
@@ -111,6 +115,10 @@
     "from_list",
     "from_numpy",
     "from_torch",
+    "hypot",
+    "index_copy",
+    "index_add",
+    "take",
     "ones",
     "strided_empty",
     "strided_from_blob",

python/infinicore/nn/functional/__init__.py

@@ -6,12 +6,13 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
-
+from .smooth_l1_loss import smooth_l1_loss
 __all__ = [
     "causal_softmax",
     "random_sample",
     "rms_norm",
     "silu",
+    "smooth_l1_loss",
     "swiglu",
     "linear",
     "embedding",

python/infinicore/nn/functional/smooth_l1_loss.py

@@ -0,0 +1,60 @@
+from typing import Optional
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+_REDUCTION_MODES = {
+    "none": 0,
+    "mean": 1,
+    "sum": 2,
+}
+
+def smooth_l1_loss(
+    input: Tensor, 
+    target: Tensor, 
+    beta: float = 1.0, 
+    reduction: str = "mean", 
+    *, 
+    out: Optional[Tensor] = None
+) -> Tensor:
+    r"""Creates a criterion that uses a squared term if the absolute
+    element-wise error falls below beta and an L1 term otherwise.
+
+    Args:
+        input (Tensor): the input tensor.
+        target (Tensor): the target tensor.
+        beta (float, optional): The threshold at which to change between L1 and L2 loss. 
+            The value must be non-negative. Default: 1.0.
+        reduction (str, optional): Specifies the reduction to apply to the output:
+            'none': no reduction will be applied,
+            'mean': the sum of the output will be divided by the number of elements in the output,
+            'sum': the output will be summed. Default: 'mean'.
+        out (Tensor, optional): the output tensor.
+
+    Returns:
+        Tensor: The loss value.
+    """
+
+    if not input.is_contiguous():
+        input = input.contiguous()
+    if not target.is_contiguous():
+        target = target.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.smooth_l1_loss_(
+            out._underlying, 
+            input._underlying, 
+            target._underlying, 
+            beta, 
+            reduction_val
+        )
+        return out
+    return Tensor(
+        _infinicore.smooth_l1_loss(
+            input._underlying, 
+            target._underlying, 
+            beta, 
+            reduction_val
+        )
+    )

python/infinicore/ops/hypot.py

@@ -0,0 +1,10 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def hypot(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.hypot(input._underlying, other._underlying))
+    _infinicore.hypot_(out._underlying, input._underlying, other._underlying)
+
+    return out