issue/1031 T1-1-17

include/infinicore/ops.hpp

@@ -3,9 +3,13 @@
 #include "ops/add.hpp"
 #include "ops/add_rms_norm.hpp"
 #include "ops/attention.hpp"
+#include "ops/avg_pool1d.hpp"
 #include "ops/causal_softmax.hpp"
+#include "ops/cross_entropy.hpp"
 #include "ops/embedding.hpp"
 #include "ops/flash_attention.hpp"
+#include "ops/hardswish.hpp"
+#include "ops/hardtanh.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/matmul.hpp"
 #include "ops/ones.hpp"

include/infinicore/ops/avg_pool1d.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class AvgPool1d {
+public:
+    using schema = void (*)(Tensor, Tensor, size_t, size_t, size_t);
+    static void execute(Tensor output, Tensor input, size_t kernel_size, size_t stride, size_t padding);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor avg_pool1d(Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0);
+void avg_pool1d_(Tensor output, Tensor input, size_t kernel_size, size_t stride = 0, size_t padding = 0);
+
+} // namespace infinicore::op

include/infinicore/ops/cross_entropy.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class CrossEntropy {
+public:
+    // Schema 定义：函数指针类型
+    // CrossEntropy 需要接收三个 Tensor: Output (Loss), Input (Logits), Target (Labels)
+    using schema = void (*)(Tensor, Tensor, Tensor);
+
+    // 执行入口
+    static void execute(Tensor output, Tensor input, Tensor target);
+
+    // 分发器访问接口
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+// ==================================================================
+// 对外 Functional API
+// ==================================================================
+
+// 1. Out-of-place 接口：
+// 输入 Logits 和 Target，内部自动创建 Output Tensor 并返回
+Tensor cross_entropy(Tensor input, Tensor target);
+
+// 2. Explicit Output 接口 (类似于 In-place 风格)：
+// 用户显式提供 Output Tensor 用于存储结果
+// 注意：虽然命名带有下划线 _，但通常 CrossEntropy 无法真正原地修改 input，
+// 所以这里只是表示“写入指定的 output 内存”
+void cross_entropy_(Tensor output, Tensor input, Tensor target);
+
+} // namespace infinicore::op

include/infinicore/ops/equal.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Equal {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor);
+
+    static void execute(Tensor out, Tensor a, Tensor b);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor equal(Tensor a, Tensor b);
+void equal_(Tensor out, Tensor a, Tensor b);
+
+} // namespace infinicore::op

include/infinicore/ops/hardswish.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class Hardswish {
+public:
+    using schema = void (*)(Tensor, Tensor);
+    static void execute(Tensor output, Tensor input);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor hardswish(Tensor input);
+void hardswish_(Tensor output, Tensor input);
+
+} // namespace infinicore::op

include/infinicore/ops/hardtanh.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+
+namespace infinicore::op {
+
+class HardTanh {
+public:
+    using schema = void (*)(Tensor, Tensor, float, float);
+    static void execute(Tensor output, Tensor input, float min_val, float max_val);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor hardtanh(Tensor input, float min_val = -1.0f, float max_val = 1.0f);
+void hardtanh_(Tensor output, Tensor input, float min_val = -1.0f, float max_val = 1.0f);
+
+} // namespace infinicore::op

include/infiniop.h

@@ -47,4 +47,10 @@
 #include "infiniop/ops/zeros.h"
 #include "infiniop/tensor_descriptor.h"
 
+#include "infiniop/ops/cross_entropy.h"
+#include "infiniop/ops/hardswish.h"
+#include "infiniop/ops/avg_pool1d.h"
+#include "infiniop/ops/equal.h"
+#include "infiniop/ops/hardtanh.h"
+
 #endif // __INFINIOP_API_H__

include/infiniop/ops/avg_pool1d.h

+#ifndef __INFINIOP_AVG_POOL1D_API_H__
+#define __INFINIOP_AVG_POOL1D_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopAvgPool1dDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateAvgPool1dDescriptor(
+    infiniopHandle_t handle,
+    infiniopAvgPool1dDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output,
+    infiniopTensorDescriptor_t input,
+    size_t kernel_size,
+    size_t stride,
+    size_t padding);
+
+__INFINI_C __export infiniStatus_t infiniopGetAvgPool1dWorkspaceSize(
+    infiniopAvgPool1dDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopAvgPool1d(
+    infiniopAvgPool1dDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyAvgPool1dDescriptor(
+    infiniopAvgPool1dDescriptor_t desc);
+
+#endif

include/infiniop/ops/cross_entropy.h

+#ifndef __INFINIOP_CROSS_ENTROPY_API_H__
+#define __INFINIOP_CROSS_ENTROPY_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopCrossEntropyDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateCrossEntropyDescriptor(
+    infiniopHandle_t handle,
+    infiniopCrossEntropyDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t y_desc,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t target_desc);
+
+__INFINI_C __export infiniStatus_t infiniopGetCrossEntropyWorkspaceSize(
+    infiniopCrossEntropyDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopCrossEntropy(
+    infiniopCrossEntropyDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *y,
+    const void *x,
+    const void *target,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyCrossEntropyDescriptor(
+    infiniopCrossEntropyDescriptor_t desc);
+
+#endif

include/infiniop/ops/equal.h

+#ifndef __INFINIOP_EQUAL_API_H__
+#define __INFINIOP_EQUAL_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopEqualDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateEqualDescriptor(
+    infiniopHandle_t handle,
+    infiniopEqualDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t c,
+    infiniopTensorDescriptor_t a,
+    infiniopTensorDescriptor_t b);
+
+__INFINI_C __export infiniStatus_t infiniopGetEqualWorkspaceSize(
+    infiniopEqualDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopEqual(
+    infiniopEqualDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *c,
+    const void *a,
+    const void *b,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyEqualDescriptor(
+    infiniopEqualDescriptor_t desc);
+
+#endif

include/infiniop/ops/hardswish.h

+#ifndef __INFINIOP_HARDSWISH_API_H__
+#define __INFINIOP_HARDSWISH_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopHardSwishDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateHardSwishDescriptor(
+    infiniopHandle_t handle,
+    infiniopHardSwishDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output,
+    infiniopTensorDescriptor_t input);
+
+__INFINI_C __export infiniStatus_t infiniopGetHardSwishWorkspaceSize(
+    infiniopHardSwishDescriptor_t desc,
+    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopHardSwish(
+    infiniopHardSwishDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyHardSwishDescriptor(
+    infiniopHardSwishDescriptor_t desc);
+
+#endif

include/infiniop/ops/hardtanh.h

+#ifndef __INFINIOP_HARDTANH_API_H__
+#define __INFINIOP_HARDTANH_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopHardTanhDescriptor_t;
+
+__INFINI_C __export infiniStatus_t infiniopCreateHardTanhDescriptor(infiniopHandle_t handle,
+                                                                    infiniopHardTanhDescriptor_t *desc_ptr,
+                                                                    infiniopTensorDescriptor_t output,
+                                                                    infiniopTensorDescriptor_t input,
+                                                                    float min_val,
+                                                                    float max_val);
+
+__INFINI_C __export infiniStatus_t infiniopGetHardTanhWorkspaceSize(infiniopHardTanhDescriptor_t desc,
+                                                                    size_t *size);
+
+__INFINI_C __export infiniStatus_t infiniopHardTanh(infiniopHardTanhDescriptor_t desc,
+                                                    void *workspace,
+                                                    size_t workspace_size,
+                                                    void *output,
+                                                    const void *input,
+                                                    void *stream);
+
+__INFINI_C __export infiniStatus_t infiniopDestroyHardTanhDescriptor(infiniopHardTanhDescriptor_t desc);
+
+#endif

python/infinicore/__init__.py

@@ -51,6 +51,8 @@ from infinicore.ops.add import add
 from infinicore.ops.add_rms_norm import add_rms_norm
 from infinicore.ops.all import all
 from infinicore.ops.attention import attention
+from infinicore.ops.cross_entropy import cross_entropy
+from infinicore.ops.equal import equal
 from infinicore.ops.kv_caching import kv_caching
 from infinicore.ops.matmul import matmul
 from infinicore.ops.mha_varlen import mha_varlen
@@ -129,11 +131,13 @@ __all__ = [
     "attention",
     "kv_caching",
     "matmul",
+    "equal",
     "mul",
     "narrow",
     "squeeze",
     "unsqueeze",
     "rearrange",
+    "cross_entropy",
     "empty",
     "empty_like",
     "from_blob",

python/infinicore/nn/functional/__init__.py

+from .avg_pool1d import avg_pool1d
 from .causal_softmax import causal_softmax
 from .embedding import embedding
 from .flash_attention import flash_attention
+from .hardswish import hardswish
+from .hardtanh import hardtanh
 from .linear import linear
 from .linear_w8a8i8 import linear_w8a8i8
 from .random_sample import random_sample
@@ -20,6 +23,9 @@ __all__ = [
     "RopeAlgo",
     "rope",
     "silu",
+    "hardswish",
+    "hardtanh",
+    "avg_pool1d",
     "swiglu",
     "linear_w8a8i8",
     "silu_and_mul",

python/infinicore/nn/functional/avg_pool1d.py

+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def avg_pool1d(
+    input: Tensor,
+    kernel_size: int,
+    stride: int | None = None,
+    padding: int = 0,
+    *,
+    out=None,
+) -> Tensor:
+    if stride is None:
+        stride = 0
+
+    if out is None:
+        return Tensor(
+            _infinicore.avg_pool1d(input._underlying, kernel_size, stride, padding)
+        )
+
+    _infinicore.avg_pool1d_(
+        out._underlying, input._underlying, kernel_size, stride, padding
+    )
+    return out

python/infinicore/nn/functional/hardswish.py

+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def hardswish(input: Tensor, inplace: bool = False, *, out=None) -> Tensor:
+    r"""Apply the Hardswish activation function element-wise."""
+
+    if (
+        infinicore.use_ntops
+        and input.device.type in ("cuda", "musa")
+        and out is None
+        and hasattr(infinicore.ntops.torch, "hardswish")
+    ):
+        try:
+            return infinicore.ntops.torch.hardswish(input, inplace=inplace)
+        except AttributeError:
+            pass
+
+    if inplace:
+        _infinicore.hardswish_(input._underlying, input._underlying)
+        return input
+
+    if out is None:
+        return Tensor(_infinicore.hardswish(input._underlying))
+
+    _infinicore.hardswish_(out._underlying, input._underlying)
+    return out

python/infinicore/nn/functional/hardtanh.py

+import infinicore
+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def hardtanh(
+    input: Tensor,
+    min_val: float = -1.0,
+    max_val: float = 1.0,
+    inplace: bool = False,
+    *,
+    out=None,
+) -> Tensor:
+    """Clamp the input tensor to the range [min_val, max_val]."""
+
+    if min_val > max_val:
+        raise ValueError("min_val must be less than or equal to max_val")
+
+    if (
+        infinicore.use_ntops
+        and input.device.type in ("cuda", "musa")
+        and out is None
+        and hasattr(infinicore.ntops.torch, "hardtanh")
+    ):
+        try:
+            return infinicore.ntops.torch.hardtanh(
+                input, min_val=min_val, max_val=max_val, inplace=inplace
+            )
+        except AttributeError:
+            pass
+
+    if inplace:
+        _infinicore.hardtanh_(
+            input._underlying, input._underlying, float(min_val), float(max_val)
+        )
+        return input
+
+    if out is None:
+        return Tensor(
+            _infinicore.hardtanh(input._underlying, float(min_val), float(max_val))
+        )
+
+    _infinicore.hardtanh_(
+        out._underlying, input._underlying, float(min_val), float(max_val)
+    )
+    return out

python/infinicore/ops/cross_entropy.py

+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def cross_entropy(
+    logits,
+    target,
+    weight=None,
+    *,
+    ignore_index=None,
+    reduction="none",
+    out=None,
+):
+    """
+    Token-wise cross entropy without reduction. The output tensor has the same
+    shape as target and uses the logits dtype.
+    """
+    if weight is not None:
+        raise NotImplementedError("class weights are not supported yet.")
+    if ignore_index is not None:
+        raise NotImplementedError("ignore_index is not supported yet.")
+    if reduction not in (None, "none"):
+        raise NotImplementedError("Only reduction='none' is implemented.")
+
+    if out is None:
+        return Tensor(_infinicore.cross_entropy(logits._underlying, target._underlying))
+
+    _infinicore.cross_entropy_(
+        out._underlying,
+        logits._underlying,
+        target._underlying,
+    )
+    return out

python/infinicore/ops/equal.py

+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def equal(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.equal(input._underlying, other._underlying))
+
+    _infinicore.equal_(out._underlying, input._underlying, other._underlying)
+    return out

python/infinicore/utils.py

-import ml_dtypes
 import numpy as np
 import torch
 
 import infinicore
 
+try:
+    import ml_dtypes
+except ModuleNotFoundError:
+    ml_dtypes = None
+
 
 def to_torch_dtype(infini_dtype):
     """Convert infinicore data type to PyTorch data type"""
@@ -57,7 +61,9 @@ def numpy_to_infinicore_dtype(numpy_dtype):
         return infinicore.float64
     elif numpy_dtype == np.float16:
         return infinicore.float16
-    elif numpy_dtype == ml_dtypes.bfloat16:
+    elif hasattr(np, "bfloat16") and numpy_dtype == np.bfloat16:
+        return infinicore.bfloat16
+    elif ml_dtypes is not None and numpy_dtype == ml_dtypes.bfloat16:
         return infinicore.bfloat16
     elif numpy_dtype == np.int8:
         return infinicore.int8
@@ -86,6 +92,13 @@ def infinicore_to_numpy_dtype(infini_dtype):
     elif infini_dtype == infinicore.int16:
         return np.int16
     elif infini_dtype == infinicore.bfloat16:
+        if hasattr(np, "bfloat16"):
+            return np.bfloat16
+        if ml_dtypes is None:
+            raise ModuleNotFoundError(
+                "ml_dtypes is required for bfloat16 numpy conversion. "
+                "Please install ml_dtypes."
+            )
         return ml_dtypes.bfloat16
     elif infini_dtype == infinicore.int32:
         return np.int32