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

include/infinicore/ops/acos.hpp

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

include/infinicore/ops/adaptive_avg_pool1d.hpp

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class AdaptiveAvgPool1d {
+public:
+    // Schema: execute(Output, Input)
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor adaptive_avg_pool1d(Tensor input, int64_t output_size);
+
+} // namespace infinicore::op

include/infinicore/ops/addbmm.hpp

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Addbmm {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, float, float);
+    static void execute(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha);
+
+    static common::OpDispatcher<schema> &dispatcher();
+};
+Tensor addbmm(Tensor input, Tensor batch1, Tensor batch2, float beta = 1.0f, float alpha = 1.0f);
+
+void addbmm_(Tensor output, Tensor input, Tensor batch1, Tensor batch2, float beta, float alpha);
+
+} // namespace infinicore::op

include/infinicore/ops/affine_grid.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <vector>
+
+namespace infinicore::op {
+
+class AffineGrid {
+public:
+    using schema = void (*)(Tensor, Tensor, bool);
+    static void execute(Tensor output, Tensor theta, bool align_corners);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+Tensor affine_grid(Tensor theta, const std::vector<int64_t>& size, bool align_corners = false);
+
+} // namespace infinicore::op

include/infinicore/ops/floor.hpp

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

include/infiniop.h

@@ -2,7 +2,11 @@
 #define __INFINIOP_API_H__
 
 #include "infiniop/handle.h"
+#include "infiniop/ops/acos.h"
 #include "infiniop/ops/add.h"
+#include "infiniop/ops/adaptive_avg_pool1d.h"
+#include "infiniop/ops/addbmm.h"
+#include "infiniop/ops/affine_grid.h"
 #include "infiniop/ops/attention.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"

include/infiniop/ops/acos.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_ACOS_API_H__
+#define __INFINIOP_ACOS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopAcosDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateAcosDescriptor(infiniopHandle_t handle,
+                                                         infiniopAcosDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t y, 
+                                                         infiniopTensorDescriptor_t x );
+
+__C __export infiniStatus_t infiniopGetAcosWorkspaceSize(infiniopAcosDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopAcos(infiniopAcosDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *y,      
+                                         const void *x, 
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyAcosDescriptor(infiniopAcosDescriptor_t desc);
+
+#endif 

include/infiniop/ops/adaptive_avg_pool1d.h

@@ -0,0 +1,34 @@
+#ifndef __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__
+#define __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__
+
+#include "../operator_descriptor.h"
+
+// 定义算子描述符类型
+typedef struct InfiniopDescriptor *infiniopAdaptiveAvgPool1dDescriptor_t;
+
+// 1. 创建算子描述符
+__C __export infiniStatus_t infiniopCreateAdaptiveAvgPool1dDescriptor(
+    infiniopHandle_t handle,
+    infiniopAdaptiveAvgPool1dDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    infiniopTensorDescriptor_t input_desc);
+
+// 2. 获取 Workspace 大小
+__C __export infiniStatus_t infiniopGetAdaptiveAvgPool1dWorkspaceSize(
+    infiniopAdaptiveAvgPool1dDescriptor_t desc, 
+    size_t *size);
+
+// 3. 执行计算
+__C __export infiniStatus_t infiniopAdaptiveAvgPool1d(
+    infiniopAdaptiveAvgPool1dDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+// 4. 销毁描述符
+__C __export infiniStatus_t infiniopDestroyAdaptiveAvgPool1dDescriptor(
+    infiniopAdaptiveAvgPool1dDescriptor_t desc);
+
+#endif // __INFINIOP_ADAPTIVE_AVG_POOL1D_API_H__

include/infiniop/ops/addbmm.h

@@ -0,0 +1,30 @@
+#ifndef __INFINIOP_ADDBMM_API_H__
+#define __INFINIOP_ADDBMM_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopAddbmmDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateAddbmmDescriptor(infiniopHandle_t handle,
+                                                           infiniopAddbmmDescriptor_t *desc_ptr,
+                                                           infiniopTensorDescriptor_t out_desc,
+                                                           infiniopTensorDescriptor_t input_desc,
+                                                           infiniopTensorDescriptor_t batch1_desc,
+                                                           infiniopTensorDescriptor_t batch2_desc,
+                                                           float alpha,
+                                                           float beta);
+
+__C __export infiniStatus_t infiniopGetAddbmmWorkspaceSize(infiniopAddbmmDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopAddbmm(infiniopAddbmmDescriptor_t desc,
+                                           void *workspace,
+                                           size_t workspace_size,
+                                           void *output,
+                                           const void *input,
+                                           const void *batch1,
+                                           const void *batch2,
+                                           void *stream);
+
+__C __export infiniStatus_t infiniopDestroyAddbmmDescriptor(infiniopAddbmmDescriptor_t desc);
+
+#endif // __INFINIOP_ADDBMM_API_H__

include/infiniop/ops/affine_grid.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_AFFINE_GRID_API_H__
+#define __INFINIOP_AFFINE_GRID_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopAffineGridDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateAffineGridDescriptor(infiniopHandle_t handle,
+                                                               infiniopAffineGridDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t output_desc,
+                                                               infiniopTensorDescriptor_t input_desc,
+                                                               uint8_t align_corners);
+
+__C __export infiniStatus_t infiniopGetAffineGridWorkspaceSize(infiniopAffineGridDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopAffineGrid(infiniopAffineGridDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *output,
+                                               const void *input,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyAffineGridDescriptor(infiniopAffineGridDescriptor_t desc);
+
+#endif // __INFINIOP_AFFINE_GRID_API_H__

include/infiniop/ops/floor.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_FLOOR_API_H__
+#define __INFINIOP_FLOOR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopFloorDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateFloorDescriptor(infiniopHandle_t handle,
+                                                         infiniopFloorDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t intput);
+
+__C __export infiniStatus_t infiniopGetFloorWorkspaceSize(infiniopFloorDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopFloor(infiniopFloorDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *intput,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyFloorDescriptor(infiniopFloorDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -40,6 +40,9 @@
     uint8,
 )
 from infinicore.ops.add import add
+from infinicore.ops.acos import acos
+from infinicore.ops.floor import floor
+from infinicore.ops.addbmm import addbmm
 from infinicore.ops.attention import attention
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
@@ -100,6 +103,9 @@
     "uint8",
     # Operations.
     "add",
+    "acos",
+    "addbmm",
+    "floor",
     "attention",
     "matmul",
     "mul",

python/infinicore/nn/functional/__init__.py

@@ -6,10 +6,13 @@
 from .rope import RopeAlgo, rope
 from .silu import silu
 from .swiglu import swiglu
-
+from .adaptive_avg_pool1d import adaptive_avg_pool1d
+from .affine_grid import affine_grid  
 __all__ = [
     "causal_softmax",
     "random_sample",
+    "adaptive_avg_pool1d",
+    "affine_grid",  
     "rms_norm",
     "silu",
     "swiglu",

python/infinicore/nn/functional/adaptive_avg_pool1d.py

@@ -0,0 +1,7 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def adaptive_avg_pool1d(input: Tensor, output_size: int) -> Tensor:
+    r"""Apply a 1D adaptive average pooling."""
+    return Tensor(_infinicore.adaptive_avg_pool1d(input._underlying, output_size))

python/infinicore/nn/functional/affine_grid.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def affine_grid(theta: Tensor, size: list[int], align_corners: bool = False) -> Tensor:
+    r"""Generates a 2D flow field (sampling grid), given a batch of affine matrices theta."""
+
+    # 直接调用底层绑定
+    # theta._underlying: 传递底层 C++ Tensor 对象
+    # size: Python list[int] 自动转换为 C++ std::vector<int64_t>
+    return Tensor(_infinicore.affine_grid(theta._underlying, size, align_corners))

python/infinicore/ops/acos.py

@@ -0,0 +1,9 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def acos(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.acos(input._underlying))
+    _infinicore.acos_(out._underlying, input._underlying)
+
+    return out

python/infinicore/ops/addbmm.py

@@ -0,0 +1,25 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+def addbmm(input, batch1, batch2, *, beta=1.0, alpha=1.0, out=None):
+    # 1. Out-of-place 模式 (如果没有指定 out)
+    if out is None:
+        return Tensor(_infinicore.addbmm(
+            input._underlying, 
+            batch1._underlying, 
+            batch2._underlying, 
+            beta, 
+            alpha
+        ))
+
+    # 2. In-place 模式 (指定了 out)
+    _infinicore.addbmm_(
+        out._underlying, 
+        input._underlying, 
+        batch1._underlying, 
+        batch2._underlying, 
+        beta, 
+        alpha
+    )
+
+    return out

python/infinicore/ops/floor.py

@@ -0,0 +1,11 @@
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def floor(input, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.floor(input._underlying))
+
+    _infinicore.floor_(out._underlying, input._underlying)
+
+    return out

src/infinicore/ops/acos/acos.cc

@@ -0,0 +1,24 @@
+#include "infinicore/ops/acos.hpp"
+
+namespace infinicore::op {
+
+common::OpDispatcher<Acos::schema> &Acos::dispatcher() {
+    static common::OpDispatcher<Acos::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void Acos::execute(Tensor output, Tensor input) {
+    dispatcher().lookup(context::getDevice().getType())(output, input);
+}
+
+Tensor acos(Tensor input) {
+    auto output = Tensor::empty(input->shape(), input->dtype(), input->device());
+    acos_(output, input);
+    return output;
+}
+
+void acos_(Tensor output, Tensor input) {
+    Acos::execute(output, input);
+}
+
+} // namespace infinicore::op