diff --git a/include/infinicore/ops/log_softmax.hpp b/include/infinicore/ops/log_softmax.hpp new file mode 100644 index 000000000..2451e81fd --- /dev/null +++ b/include/infinicore/ops/log_softmax.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogSoftmax { +public: + // Schema signature: output(out), input, dim + using schema = void (*)(Tensor, Tensor, int64_t); + + static void execute(Tensor output, Tensor input, int64_t dim); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +Tensor log_softmax(Tensor input, int64_t dim); + +// In-place/Output-provided API +void log_softmax_(Tensor output, Tensor input, int64_t dim); + +} // namespace infinicore::op \ No newline at end of file diff --git a/include/infinicore/ops/logaddexp.hpp b/include/infinicore/ops/logaddexp.hpp new file mode 100644 index 000000000..197918d52 --- /dev/null +++ b/include/infinicore/ops/logaddexp.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogAddExp { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor logaddexp(Tensor a, Tensor b); +void logaddexp_(Tensor c, Tensor a, Tensor b); + +} // namespace infinicore::op \ No newline at end of file diff --git a/include/infinicore/ops/logaddexp2.hpp b/include/infinicore/ops/logaddexp2.hpp new file mode 100644 index 000000000..62fe7fd14 --- /dev/null +++ b/include/infinicore/ops/logaddexp2.hpp @@ -0,0 +1,18 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class LogAddExp2 { +public: + using schema = void (*)(Tensor, Tensor, Tensor); + static void execute(Tensor c, Tensor a, Tensor b); + static common::OpDispatcher &dispatcher(); +}; + +Tensor logaddexp2(Tensor a, Tensor b); +void logaddexp2_(Tensor c, Tensor a, Tensor b); + +} // namespace infinicore::op \ No newline at end of file diff --git a/include/infinicore/ops/triplet_margin_with_distance_loss.hpp b/include/infinicore/ops/triplet_margin_with_distance_loss.hpp new file mode 100644 index 000000000..1886b8a02 --- /dev/null +++ b/include/infinicore/ops/triplet_margin_with_distance_loss.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +class TripletMarginWithDistanceLoss { +public: + // Schema signature: output(out), anchor, positive, negative, margin, swap, reduction + using schema = void (*)(Tensor, Tensor, Tensor, Tensor, double, bool, int64_t); + + static void execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +// margin default 1.0, swap default false, reduction default 1 (Mean) typically +Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin = 1.0, bool swap = false, int64_t reduction = 1); + +// In-place/Output-provided API +void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction); + +} // namespace infinicore::op \ No newline at end of file diff --git a/include/infinicore/ops/upsample_nearest.hpp b/include/infinicore/ops/upsample_nearest.hpp new file mode 100644 index 000000000..51534ab51 --- /dev/null +++ b/include/infinicore/ops/upsample_nearest.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include "../device.hpp" +#include "common/op.hpp" +#include + +namespace infinicore::op { + +class UpsampleNearest { +public: + // Schema signature: output(out), input + // Note: Scales are inferred from output.shape / input.shape + using schema = void (*)(Tensor, Tensor); + + static void execute(Tensor output, Tensor input); + static common::OpDispatcher &dispatcher(); +}; + +// Functional API: Returns the result tensor +// Requires output_size to calculate the shape of the result tensor +Tensor upsample_nearest(Tensor input, const std::vector& output_size); + +// In-place/Output-provided API +void upsample_nearest_(Tensor output, Tensor input); + +} // namespace infinicore::op \ No newline at end of file diff --git a/include/infiniop.h b/include/infiniop.h index c0a09fcb4..1b5dfa414 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -12,8 +12,12 @@ #include "infiniop/ops/gelu.h" #include "infiniop/ops/gemm.h" #include "infiniop/ops/layer_norm.h" -#include "infiniop/ops/logsoftmax.h" #include "infiniop/ops/lp_norm.h" +#include "infiniop/ops/log_softmax.h" +#include "infiniop/ops/upsample_nearest.h" +#include "infiniop/ops/triplet_margin_with_distance_loss.h" +#include "infiniop/ops/logaddexp.h" +#include "infiniop/ops/logaddexp2.h" #include "infiniop/ops/mul.h" #include "infiniop/ops/ones.h" #include "infiniop/ops/paged_attention.h" diff --git a/include/infiniop/ops/logsoftmax.h b/include/infiniop/ops/log_softmax.h similarity index 52% rename from include/infiniop/ops/logsoftmax.h rename to include/infiniop/ops/log_softmax.h index 1b944424c..ebaf73dcf 100644 --- a/include/infiniop/ops/logsoftmax.h +++ b/include/infiniop/ops/log_softmax.h @@ -1,24 +1,25 @@ -#ifndef __INFINIOP_LOGSOFTMAX_API_H__ -#define __INFINIOP_LOGSOFTMAX_API_H__ +#ifndef __INFINIOP_LOG_SOFTMAX_API_H__ +#define __INFINIOP_LOG_SOFTMAX_API_H__ #include "../operator_descriptor.h" typedef struct InfiniopDescriptor *infiniopLogSoftmaxDescriptor_t; __C __export infiniStatus_t infiniopCreateLogSoftmaxDescriptor(infiniopHandle_t handle, - infiniopLogSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc); + infiniopLogSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int dim); __C __export infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size); __C __export infiniStatus_t infiniopLogSoftmax(infiniopLogSoftmaxDescriptor_t desc, void *workspace, size_t workspace_size, - void *y, - const void *x, + void *output, + const void *input, void *stream); __C __export infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc); -#endif +#endif // __INFINIOP_LOG_SOFTMAX_API_H__ \ No newline at end of file diff --git a/include/infiniop/ops/logaddexp.h b/include/infiniop/ops/logaddexp.h new file mode 100644 index 000000000..6e6955598 --- /dev/null +++ b/include/infiniop/ops/logaddexp.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_LOGADDEXP_API_H__ +#define __INFINIOP_LOGADDEXP_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogAddExpDescriptor_t; + +__C __export infiniStatus_t infiniopCreateLogAddExpDescriptor(infiniopHandle_t handle, + infiniopLogAddExpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__C __export infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size); + +__C __export infiniStatus_t infiniopLogAddExp(infiniopLogAddExpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__C __export infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDescriptor_t desc); + +#endif // __INFINIOP_LOGADDEXP_API_H__ \ No newline at end of file diff --git a/include/infiniop/ops/logaddexp2.h b/include/infiniop/ops/logaddexp2.h new file mode 100644 index 000000000..ddf5ea530 --- /dev/null +++ b/include/infiniop/ops/logaddexp2.h @@ -0,0 +1,26 @@ +#ifndef __INFINIOP_LOGADDEXP2_API_H__ +#define __INFINIOP_LOGADDEXP2_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopLogAddExp2Descriptor_t; + +__C __export infiniStatus_t infiniopCreateLogAddExp2Descriptor(infiniopHandle_t handle, + infiniopLogAddExp2Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b); + +__C __export infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size); + +__C __export infiniStatus_t infiniopLogAddExp2(infiniopLogAddExp2Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream); + +__C __export infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2Descriptor_t desc); + +#endif // __INFINIOP_LOGADDEXP2_API_H__ \ No newline at end of file diff --git a/include/infiniop/ops/triplet_margin_with_distance_loss.h b/include/infiniop/ops/triplet_margin_with_distance_loss.h new file mode 100644 index 000000000..262cdfd18 --- /dev/null +++ b/include/infiniop/ops/triplet_margin_with_distance_loss.h @@ -0,0 +1,32 @@ +#ifndef __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ +#define __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopTripletMarginWithDistanceLossDescriptor_t; + +__C __export infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor( + infiniopHandle_t handle, + infiniopTripletMarginWithDistanceLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t anchor, + infiniopTensorDescriptor_t positive, + infiniopTensorDescriptor_t negative, + float margin, + int swap, + int reduction); +__C __export infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize( + infiniopTripletMarginWithDistanceLossDescriptor_t desc, + size_t *size); +__C __export infiniStatus_t infiniopTripletMarginWithDistanceLoss(infiniopTripletMarginWithDistanceLossDescriptor_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 infiniopDestroyTripletMarginWithDistanceLossDescriptor( + infiniopTripletMarginWithDistanceLossDescriptor_t desc); +#endif // __INFINIOP_TRIPLET_MARGIN_WITH_DISTANCE_LOSS_API_H__ \ No newline at end of file diff --git a/include/infiniop/ops/upsample_nearest.h b/include/infiniop/ops/upsample_nearest.h new file mode 100644 index 000000000..f81d6004a --- /dev/null +++ b/include/infiniop/ops/upsample_nearest.h @@ -0,0 +1,24 @@ +#ifndef __INFINIOP_UPSAMPLE_NEAREST_API_H__ +#define __INFINIOP_UPSAMPLE_NEAREST_API_H__ + +#include "../operator_descriptor.h" + +typedef struct InfiniopDescriptor *infiniopUpsampleNearestDescriptor_t; + +__C __export infiniStatus_t infiniopCreateUpsampleNearestDescriptor(infiniopHandle_t handle, + infiniopUpsampleNearestDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input); + +__C __export infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size); + +__C __export infiniStatus_t infiniopUpsampleNearest(infiniopUpsampleNearestDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream); + +__C __export infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc); + +#endif // __INFINIOP_UPSAMPLE_NEAREST_API_H__ \ No newline at end of file diff --git a/python/infinicore/__init__.py b/python/infinicore/__init__.py index b7288f3ac..381a3ad59 100644 --- a/python/infinicore/__init__.py +++ b/python/infinicore/__init__.py @@ -42,6 +42,8 @@ from infinicore.ops.add import add from infinicore.ops.add_rms_norm import add_rms_norm, add_rms_norm_ from infinicore.ops.attention import attention +from infinicore.ops.logaddexp2 import logaddexp2 +from infinicore.ops.logaddexp import logaddexp from infinicore.ops.matmul import matmul from infinicore.ops.mul import mul from infinicore.ops.narrow import narrow @@ -109,6 +111,8 @@ "add_rms_norm", "add_rms_norm_", "attention", + "logaddexp", + "logaddexp2", "matmul", "mul", "narrow", diff --git a/python/infinicore/nn/functional/__init__.py b/python/infinicore/nn/functional/__init__.py index 255079790..10b3368f5 100644 --- a/python/infinicore/nn/functional/__init__.py +++ b/python/infinicore/nn/functional/__init__.py @@ -1,19 +1,25 @@ from .causal_softmax import causal_softmax from .embedding import embedding from .linear import linear +from .log_softmax import log_softmax from .random_sample import random_sample from .rms_norm import rms_norm from .rope import RopeAlgo, rope from .silu import silu from .swiglu import swiglu - +from .triplet_margin_with_distance_loss import triplet_margin_with_distance_loss +from .upsample_nearest import upsample_nearest, interpolate __all__ = [ "causal_softmax", "random_sample", "rms_norm", "silu", "swiglu", + "interpolate", "linear", + "log_softmax", + "upsample_nearest", + "triplet_margin_with_distance_loss", "embedding", "rope", "RopeAlgo", diff --git a/python/infinicore/nn/functional/log_softmax.py b/python/infinicore/nn/functional/log_softmax.py new file mode 100644 index 000000000..373b98748 --- /dev/null +++ b/python/infinicore/nn/functional/log_softmax.py @@ -0,0 +1,36 @@ +from typing import Optional +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + +def log_softmax( + input: Tensor, + dim: int, + *, + out: Optional[Tensor] = None +) -> Tensor: + r"""Applies a softmax followed by a logarithm. + While mathematically equivalent to log(softmax(x)), doing these two + operations separately is slower and numerically unstable. This function + uses an alternative formulation to compute the output and gradient correctly. + """ + + if not input.is_contiguous(): + input = input.contiguous() + + if out is not None: + if not isinstance(out, Tensor): + raise ValueError("out must be a Tensor") + + _infinicore.log_softmax_( + out._underlying, + input._underlying, + dim + ) + return out + + ret = _infinicore.log_softmax( + input._underlying, + dim + ) + + return Tensor(ret) \ No newline at end of file diff --git a/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py b/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py new file mode 100644 index 000000000..778a51825 --- /dev/null +++ b/python/infinicore/nn/functional/triplet_margin_with_distance_loss.py @@ -0,0 +1,56 @@ +from typing import Optional, Union +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + +def triplet_margin_with_distance_loss( + anchor: Tensor, + positive: Tensor, + negative: Tensor, + *, + margin: float = 1.0, + swap: bool = False, + reduction: str = "mean", + out: Optional[Tensor] = None +) -> Tensor: + r"""Calculates the triplet margin loss for a given triplet of tensors. + The loss is defined as: L(a, p, n) = max(d(a, p) - d(a, n) + margin, 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() + + reduction_map = {"none": 0, "mean": 1, "sum": 2} + if reduction not in reduction_map: + raise ValueError(f"Invalid reduction mode: {reduction}") + + reduction_val = reduction_map[reduction] + + if out is not None: + if not isinstance(out, Tensor): + raise ValueError("out must be a Tensor") + + _infinicore.triplet_margin_with_distance_loss_( + out._underlying, + anchor._underlying, + positive._underlying, + negative._underlying, + margin, + swap, + reduction_val + ) + return out + + ret = _infinicore.triplet_margin_with_distance_loss( + anchor._underlying, + positive._underlying, + negative._underlying, + margin, + swap, + reduction_val + ) + + return Tensor(ret) \ No newline at end of file diff --git a/python/infinicore/nn/functional/upsample_nearest.py b/python/infinicore/nn/functional/upsample_nearest.py new file mode 100644 index 000000000..13cf847a3 --- /dev/null +++ b/python/infinicore/nn/functional/upsample_nearest.py @@ -0,0 +1,166 @@ +from typing import Optional, Union, Sequence +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def upsample_nearest( + input: Tensor, + size: Optional[Union[int, Sequence[int]]] = None, + scale_factor: Optional[Union[float, Sequence[float]]] = None, + *, + out: Optional[Tensor] = None +) -> Tensor: + if not input.is_contiguous(): + input = input.contiguous() + + if (size is None) == (scale_factor is None): + raise ValueError("Either size or scale_factor should be defined, but not both.") + + ndim = len(input.shape) + output_size = [] + + if size is not None: + if isinstance(size, int): + if ndim == 3: + output_size = [size] + else: + output_size = [size, size] + elif isinstance(size, (list, tuple)): + output_size = [int(s) for s in size] + else: + raise ValueError("size must be int or sequence of int") + else: + if isinstance(scale_factor, (float, int)): + scales = [float(scale_factor)] + elif isinstance(scale_factor, (list, tuple)): + scales = [float(s) for s in scale_factor] + else: + raise ValueError("scale_factor must be float or sequence of float") + + if ndim == 3: + w_in = input.shape[-1] + scale_w = scales[0] if len(scales) == 1 else scales[-1] + output_size = [int(w_in * scale_w)] + else: + if len(scales) == 1: + scale_h = scale_w = scales[0] + elif len(scales) >= 2: + scale_h, scale_w = scales[0], scales[1] + else: + raise ValueError("scale_factor sequence length mismatch") + + h_in = input.shape[-2] + w_in = input.shape[-1] + output_size = [int(h_in * scale_h), int(w_in * scale_w)] + + if out is not None: + if not out.is_contiguous(): + raise RuntimeError("out tensor must be contiguous") + + _infinicore.upsample_nearest_( + out._underlying, + input._underlying + ) + return out + + return Tensor( + _infinicore.upsample_nearest( + input._underlying, + output_size + ) + ) + + +def upsample_bilinear( + input: Tensor, + size: Optional[Union[int, Sequence[int]]] = None, + scale_factor: Optional[Union[float, Sequence[float]]] = None, + align_corners: bool = False, + *, + out: Optional[Tensor] = None +) -> Tensor: + if not input.is_contiguous(): + input = input.contiguous() + + if (size is None) == (scale_factor is None): + raise ValueError("Either size or scale_factor should be defined, but not both.") + + ndim = len(input.shape) + output_size = [] + + if size is not None: + if isinstance(size, int): + if ndim == 3: + output_size = [size] + else: + output_size = [size, size] + elif isinstance(size, (list, tuple)): + output_size = [int(s) for s in size] + else: + raise ValueError("size must be int or sequence of int") + else: + if isinstance(scale_factor, (float, int)): + scales = [float(scale_factor)] + elif isinstance(scale_factor, (list, tuple)): + scales = [float(s) for s in scale_factor] + else: + raise ValueError("scale_factor must be float or sequence of float") + + if ndim == 3: + w_in = input.shape[-1] + scale_w = scales[0] if len(scales) == 1 else scales[-1] + output_size = [int(w_in * scale_w)] + else: + if len(scales) == 1: + scale_h = scale_w = scales[0] + elif len(scales) >= 2: + scale_h, scale_w = scales[0], scales[1] + else: + raise ValueError("scale_factor sequence length mismatch") + + h_in = input.shape[-2] + w_in = input.shape[-1] + output_size = [int(h_in * scale_h), int(w_in * scale_w)] + + if out is not None: + if not out.is_contiguous(): + raise RuntimeError("out tensor must be contiguous") + + _infinicore.upsample_bilinear_( + out._underlying, + input._underlying, + align_corners + ) + return out + + return Tensor( + _infinicore.upsample_bilinear( + input._underlying, + output_size, + align_corners + ) + ) + + +def interpolate( + input: Tensor, + size: Optional[Union[int, Sequence[int]]] = None, + scale_factor: Optional[Union[float, Sequence[float]]] = None, + mode: str = 'nearest', + align_corners: Optional[bool] = None, + recompute_scale_factor: Optional[bool] = None +) -> Tensor: + if mode == 'nearest': + if align_corners is not None: + raise ValueError( + "align_corners option can only be set with the " + "interpolating modes: linear | bilinear | bicubic | trilinear" + ) + return upsample_nearest(input, size, scale_factor) + + if mode == 'bilinear': + if align_corners is None: + align_corners = False + return upsample_bilinear(input, size, scale_factor, align_corners) + + raise NotImplementedError(f"Interpolation mode '{mode}' is not currently supported.") diff --git a/python/infinicore/ops/logaddexp.py b/python/infinicore/ops/logaddexp.py new file mode 100644 index 000000000..c2cd26d3f --- /dev/null +++ b/python/infinicore/ops/logaddexp.py @@ -0,0 +1,11 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def logaddexp(input, other, *, out=None): + if out is None: + return Tensor(_infinicore.logaddexp(input._underlying, other._underlying)) + + _infinicore.logaddexp_(out._underlying, input._underlying, other._underlying) + + return out \ No newline at end of file diff --git a/python/infinicore/ops/logaddexp2.py b/python/infinicore/ops/logaddexp2.py new file mode 100644 index 000000000..65ffef7fe --- /dev/null +++ b/python/infinicore/ops/logaddexp2.py @@ -0,0 +1,11 @@ +from infinicore.lib import _infinicore +from infinicore.tensor import Tensor + + +def logaddexp2(input, other, *, out=None): + if out is None: + return Tensor(_infinicore.logaddexp2(input._underlying, other._underlying)) + + _infinicore.logaddexp2_(out._underlying, input._underlying, other._underlying) + + return out \ No newline at end of file diff --git a/src/infinicore/ops/log_softmax/log_softmax.cc b/src/infinicore/ops/log_softmax/log_softmax.cc new file mode 100644 index 000000000..2b2c24530 --- /dev/null +++ b/src/infinicore/ops/log_softmax/log_softmax.cc @@ -0,0 +1,34 @@ +#include "infinicore/ops/log_softmax.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &LogSoftmax::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogSoftmax::execute(Tensor output, Tensor input, int64_t dim) { + dispatcher().lookup(context::getDevice().getType())(output, input, dim); +} + +// 3. 函数式接口 +Tensor log_softmax(Tensor input, int64_t dim) { + int64_t ndim = input->shape().size(); + + // 处理负数维度 + if (dim < 0) { + dim += ndim; + } + + // LogSoftmax 输出形状与输入一致，dtype 与 input 一致 + auto output = Tensor::empty(input->shape(), input->dtype(), input->device()); + log_softmax_(output, input, dim); + return output; +} + +void log_softmax_(Tensor output, Tensor input, int64_t dim) { + LogSoftmax::execute(output, input, dim); +} + +} // namespace infinicore::op \ No newline at end of file diff --git a/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc b/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc new file mode 100644 index 000000000..5629551d8 --- /dev/null +++ b/src/infinicore/ops/log_softmax/log_softmax_infiniop.cc @@ -0,0 +1,65 @@ +#include "../../utils.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/common/cache.hpp" +#include "infinicore/ops/log_softmax.hpp" +#include + +namespace infinicore::op::log_softmax_impl::infiniop { + +// 定义描述符缓存 +thread_local common::OpCache caches( + 100, // capacity + [](infiniopLogSoftmaxDescriptor_t &desc) { + if (desc != nullptr) { + INFINICORE_CHECK_ERROR(infiniopDestroyLogSoftmaxDescriptor(desc)); + desc = nullptr; + } + }); + +void calculate(Tensor output, Tensor input, int64_t dim) { + size_t seed = hash_combine(output, input, dim); + + auto device_type = context::getDevice().getType(); + auto device_index = context::getDevice().getIndex(); + + auto &cache = caches.getCache(device_type, device_index); + + auto desc_opt = cache.get(seed); + infiniopLogSoftmaxDescriptor_t desc = nullptr; + + if (!desc_opt) { + // 3. 创建描述符 + INFINICORE_CHECK_ERROR(infiniopCreateLogSoftmaxDescriptor( + context::getInfiniopHandle(input->device()), + &desc, + output->desc(), + input->desc(), + static_cast(dim) + )); + + cache.put(seed, desc); + } else { + desc = *desc_opt; + } + + // 4. 获取 Workspace 并执行 + size_t workspace_size = 0; + INFINICORE_CHECK_ERROR(infiniopGetLogSoftmaxWorkspaceSize(desc, &workspace_size)); + std::shared_ptr workspace = context::allocateMemory(workspace_size); + + INFINICORE_CHECK_ERROR(infiniopLogSoftmax( + desc, + workspace->data(), + workspace_size, + output->data(), + input->data(), + context::getStream() + )); +} + +static bool registered = []() { + LogSoftmax::dispatcher().registerAll(&calculate, false); + return true; +}(); + +} // namespace infinicore::op::log_softmax_impl::infiniop \ No newline at end of file diff --git a/src/infinicore/ops/logaddexp/logaddexp.cc b/src/infinicore/ops/logaddexp/logaddexp.cc new file mode 100644 index 000000000..5481d6f0b --- /dev/null +++ b/src/infinicore/ops/logaddexp/logaddexp.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/logaddexp.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &LogAddExp::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogAddExp::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().getType())(c, a, b); +} + +Tensor logaddexp(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + logaddexp_(c, a, b); + return c; +} + +void logaddexp_(Tensor c, Tensor a, Tensor b) { + LogAddExp::execute(c, a, b); +} + +} // namespace infinicore::op \ No newline at end of file diff --git a/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc b/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc new file mode 100644 index 000000000..601458924 --- /dev/null +++ b/src/infinicore/ops/logaddexp/logaddexp_infiniop.cc @@ -0,0 +1,48 @@ +#include "../../utils.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/logaddexp.hpp" +#include "infinicore/ops/common/cache.hpp" +#include + +namespace infinicore::op::logaddexp_impl::infiniop { +thread_local common::OpCache caches( + 100, // capacity + [](infiniopLogAddExpDescriptor_t &desc) { + if (desc != nullptr) { + INFINICORE_CHECK_ERROR(infiniopDestroyLogAddExpDescriptor(desc)); + desc = nullptr; + } + }); + +void calculate(Tensor c, Tensor a, Tensor b) { + size_t seed = hash_combine(c, a, b); + + auto device = context::getDevice(); + auto &cache = caches.getCache(device); + + auto desc_opt = cache.get(seed); + infiniopLogAddExpDescriptor_t desc = nullptr; + + if (!desc_opt) { + INFINICORE_CHECK_ERROR(infiniopCreateLogAddExpDescriptor( + context::getInfiniopHandle(device), &desc, + c->desc(), a->desc(), b->desc())); + cache.put(seed, desc); + } else { + desc = *desc_opt; + } + size_t workspace_size = 0; + INFINICORE_CHECK_ERROR(infiniopGetLogAddExpWorkspaceSize(desc, &workspace_size)); + std::shared_ptr workspace = context::allocateMemory(workspace_size); + + INFINICORE_CHECK_ERROR(infiniopLogAddExp( + desc, workspace->data(), workspace_size, + c->data(), a->data(), b->data(), context::getStream())); +} + +static bool registered = []() { + LogAddExp::dispatcher().registerAll(&calculate, false); + return true; +}(); + +} // namespace infinicore::op::logaddexp_impl::infiniop \ No newline at end of file diff --git a/src/infinicore/ops/logaddexp2/logaddxep2.cc b/src/infinicore/ops/logaddexp2/logaddxep2.cc new file mode 100644 index 000000000..4dfc97839 --- /dev/null +++ b/src/infinicore/ops/logaddexp2/logaddxep2.cc @@ -0,0 +1,27 @@ +#include "infinicore/ops/logaddexp2.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +common::OpDispatcher &LogAddExp2::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void LogAddExp2::execute(Tensor c, Tensor a, Tensor b) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b); + infinicore::context::setDevice(c->device()); + dispatcher().lookup(c->device().getType())(c, a, b); +} + +Tensor logaddexp2(Tensor a, Tensor b) { + auto c = Tensor::empty(a->shape(), a->dtype(), a->device()); + logaddexp2_(c, a, b); + return c; +} + +void logaddexp2_(Tensor c, Tensor a, Tensor b) { + LogAddExp2::execute(c, a, b); +} + +} // namespace infinicore::op \ No newline at end of file diff --git a/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc b/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc new file mode 100644 index 000000000..690c41230 --- /dev/null +++ b/src/infinicore/ops/logaddexp2/logaddxep2_infiniop.cc @@ -0,0 +1,48 @@ +#include "../../utils.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/logaddexp2.hpp" +#include "infinicore/ops/common/cache.hpp" +#include + +namespace infinicore::op::logaddexp2_impl::infiniop { +thread_local common::OpCache caches( + 100, // capacity + [](infiniopLogAddExp2Descriptor_t &desc) { + if (desc != nullptr) { + INFINICORE_CHECK_ERROR(infiniopDestroyLogAddExp2Descriptor(desc)); + desc = nullptr; + } + }); + +void calculate(Tensor c, Tensor a, Tensor b) { + size_t seed = hash_combine(c, a, b); + + auto device = context::getDevice(); + auto &cache = caches.getCache(device); + + auto desc_opt = cache.get(seed); + infiniopLogAddExp2Descriptor_t desc = nullptr; + + if (!desc_opt) { + INFINICORE_CHECK_ERROR(infiniopCreateLogAddExp2Descriptor( + context::getInfiniopHandle(device), &desc, + c->desc(), a->desc(), b->desc())); + cache.put(seed, desc); + } else { + desc = *desc_opt; + } + size_t workspace_size = 0; + INFINICORE_CHECK_ERROR(infiniopGetLogAddExp2WorkspaceSize(desc, &workspace_size)); + std::shared_ptr workspace = context::allocateMemory(workspace_size); + + INFINICORE_CHECK_ERROR(infiniopLogAddExp2( + desc, workspace->data(), workspace_size, + c->data(), a->data(), b->data(), context::getStream())); +} + +static bool registered = []() { + LogAddExp2::dispatcher().registerAll(&calculate, false); + return true; +}(); + +} // namespace infinicore::op::logaddexp2_impl::infiniop \ No newline at end of file diff --git a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc new file mode 100644 index 000000000..d1c0b8544 --- /dev/null +++ b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.cc @@ -0,0 +1,38 @@ +#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" + +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &TripletMarginWithDistanceLoss::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void TripletMarginWithDistanceLoss::execute(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + dispatcher().lookup(context::getDevice().getType())(output, anchor, positive, negative, margin, swap, reduction); +} + +// 3. 函数式接口 +Tensor triplet_margin_with_distance_loss(Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + Shape out_shape; + + // 推断输出形状 + if (reduction == 0) { + // Reduction::None -> 输出形状取决于输入的广播结果 + out_shape = anchor->shape(); + } else { + // Reduction::Mean 或 Reduction::Sum -> 输出为标量 + out_shape = {}; + } + + auto output = Tensor::empty(out_shape, anchor->dtype(), anchor->device()); + + triplet_margin_with_distance_loss_(output, anchor, positive, negative, margin, swap, reduction); + return output; +} + +void triplet_margin_with_distance_loss_(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + TripletMarginWithDistanceLoss::execute(output, anchor, positive, negative, margin, swap, reduction); +} + +} // namespace infinicore::op \ No newline at end of file diff --git a/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc new file mode 100644 index 000000000..f0b5ea402 --- /dev/null +++ b/src/infinicore/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss_infiniop.cc @@ -0,0 +1,68 @@ +#include "../../utils.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/common/cache.hpp" +#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" +#include + +namespace infinicore::op::triplet_margin_with_distance_loss_impl::infiniop { + +// 定义描述符缓存 +thread_local common::OpCache caches( + 100, // capacity + [](infiniopTripletMarginWithDistanceLossDescriptor_t &desc) { + if (desc != nullptr) { + INFINICORE_CHECK_ERROR(infiniopDestroyTripletMarginWithDistanceLossDescriptor(desc)); + desc = nullptr; + } + }); + +void calculate(Tensor output, Tensor anchor, Tensor positive, Tensor negative, double margin, bool swap, int64_t reduction) { + size_t seed = hash_combine(output, anchor, positive, negative, margin, swap, reduction); + + auto device_type = context::getDevice().getType(); + auto device_index = context::getDevice().getIndex(); + + auto &cache = caches.getCache(device_type, device_index); + + auto desc_opt = cache.get(seed); + infiniopTripletMarginWithDistanceLossDescriptor_t desc = nullptr; + + if (!desc_opt) { + INFINICORE_CHECK_ERROR(infiniopCreateTripletMarginWithDistanceLossDescriptor( + context::getInfiniopHandle(anchor->device()), + &desc, + output->desc(), + anchor->desc(), + positive->desc(), + negative->desc(), + static_cast(margin), + static_cast(swap), + static_cast(reduction) + )); + + cache.put(seed, desc); + } else { + desc = *desc_opt; + } + size_t workspace_size = 0; + INFINICORE_CHECK_ERROR(infiniopGetTripletMarginWithDistanceLossWorkspaceSize(desc, &workspace_size)); + std::shared_ptr workspace = context::allocateMemory(workspace_size); + + INFINICORE_CHECK_ERROR(infiniopTripletMarginWithDistanceLoss( + desc, + workspace->data(), + workspace_size, + output->data(), + anchor->data(), + positive->data(), + negative->data(), + context::getStream() + )); +} + +static bool registered = []() { + TripletMarginWithDistanceLoss::dispatcher().registerAll(&calculate, false); + return true; +}(); + +} // namespace infinicore::op::triplet_margin_with_distance_loss_impl::infiniop \ No newline at end of file diff --git a/src/infinicore/ops/upsample_nearest/upsample_nearest.cc b/src/infinicore/ops/upsample_nearest/upsample_nearest.cc new file mode 100644 index 000000000..42aa8af06 --- /dev/null +++ b/src/infinicore/ops/upsample_nearest/upsample_nearest.cc @@ -0,0 +1,61 @@ +#include "infinicore/ops/upsample_nearest.hpp" +#include +namespace infinicore::op { + +// 1. 定义 Dispatcher 单例 +common::OpDispatcher &UpsampleNearest::dispatcher() { + static common::OpDispatcher dispatcher_; + return dispatcher_; +}; + +void UpsampleNearest::execute(Tensor output, Tensor input) { + dispatcher().lookup(context::getDevice().getType())(output, input); +} + +// 3. 函数式接口 +Tensor upsample_nearest(Tensor input, const std::vector& output_size) { + Shape input_shape = input->shape(); + size_t ndim = input_shape.size(); + + // 校验 + if (ndim < 3 || ndim > 4) { + if (ndim != 3 && ndim != 4) { + throw std::runtime_error("upsample_nearest: Only supports 3D (N,C,W) or 4D (N,C,H,W) input"); + } + } + + Shape output_shape = input_shape; + + if (ndim == 3) { + // [N, C, W] + // output_size 可能是 [W_out] (size=1) 或者 [1, W_out] (size=2) + int64_t target_w = 0; + if (output_size.size() == 1) { + target_w = output_size[0]; + } else if (output_size.size() == 2) { + target_w = output_size[1]; + } else { + throw std::runtime_error("upsample_nearest: output_size for 3D input must be [w] or [1, w]"); + } + output_shape[2] = target_w; + + } else if (ndim == 4) { + // [N, C, H, W] + if (output_size.size() != 2) { + throw std::runtime_error("upsample_nearest: output_size for 4D input must be [h, w]"); + } + output_shape[2] = output_size[0]; + output_shape[3] = output_size[1]; + } + + auto output = Tensor::empty(output_shape, input->dtype(), input->device()); + + upsample_nearest_(output, input); + return output; +} + +void upsample_nearest_(Tensor output, Tensor input) { + UpsampleNearest::execute(output, input); +} + +} // namespace infinicore::op \ No newline at end of file diff --git a/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc b/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc new file mode 100644 index 000000000..3c4e327e7 --- /dev/null +++ b/src/infinicore/ops/upsample_nearest/upsample_nearest_infiniop.cc @@ -0,0 +1,61 @@ +#include "../../utils.hpp" +#include "infinicore/common/hash.hpp" +#include "infinicore/ops/common/cache.hpp" +#include "infinicore/ops/upsample_nearest.hpp" +#include + +namespace infinicore::op::upsample_nearest_impl::infiniop { + +// 定义描述符缓存 +thread_local common::OpCache caches( + 100, // capacity + [](infiniopUpsampleNearestDescriptor_t &desc) { + if (desc != nullptr) { + INFINICORE_CHECK_ERROR(infiniopDestroyUpsampleNearestDescriptor(desc)); + desc = nullptr; + } + }); + +void calculate(Tensor output, Tensor input) { + size_t seed = hash_combine(output, input); + + auto device_type = context::getDevice().getType(); + auto device_index = context::getDevice().getIndex(); + + auto &cache = caches.getCache(device_type, device_index); + + auto desc_opt = cache.get(seed); + infiniopUpsampleNearestDescriptor_t desc = nullptr; + + if (!desc_opt) { + INFINICORE_CHECK_ERROR(infiniopCreateUpsampleNearestDescriptor( + context::getInfiniopHandle(output->device()), + &desc, + output->desc(), + input->desc() + )); + + cache.put(seed, desc); + } else { + desc = *desc_opt; + } + size_t workspace_size = 0; + INFINICORE_CHECK_ERROR(infiniopGetUpsampleNearestWorkspaceSize(desc, &workspace_size)); + std::shared_ptr workspace = context::allocateMemory(workspace_size); + + INFINICORE_CHECK_ERROR(infiniopUpsampleNearest( + desc, + workspace->data(), + workspace_size, + output->data(), + input->data(), + context::getStream() + )); +} + +static bool registered = []() { + UpsampleNearest::dispatcher().registerAll(&calculate, false); + return true; +}(); + +} // namespace infinicore::op::upsample_nearest_impl::infiniop \ No newline at end of file diff --git a/src/infinicore/pybind11/ops.hpp b/src/infinicore/pybind11/ops.hpp index 431c3a37b..d1ac0e198 100644 --- a/src/infinicore/pybind11/ops.hpp +++ b/src/infinicore/pybind11/ops.hpp @@ -8,6 +8,9 @@ #include "ops/causal_softmax.hpp" #include "ops/embedding.hpp" #include "ops/linear.hpp" +#include "ops/log_softmax.hpp" +#include "ops/logaddexp.hpp" +#include "ops/logaddexp2.hpp" #include "ops/matmul.hpp" #include "ops/mul.hpp" #include "ops/paged_attention.hpp" @@ -18,7 +21,8 @@ #include "ops/rope.hpp" #include "ops/silu.hpp" #include "ops/swiglu.hpp" - +#include "ops/triplet_margin_with_distance_loss.hpp" +#include "ops/upsample_nearest.hpp" namespace py = pybind11; namespace infinicore::ops { @@ -30,6 +34,9 @@ inline void bind(py::module &m) { bind_causal_softmax(m); bind_random_sample(m); bind_linear(m); + bind_log_softmax(m); + bind_logaddexp(m); + bind_logaddexp2(m); bind_matmul(m); bind_mul(m); bind_paged_attention(m); @@ -39,6 +46,8 @@ inline void bind(py::module &m) { bind_silu(m); bind_swiglu(m); bind_rope(m); + bind_triplet_margin_with_distance_loss(m); + bind_upsample_nearest(m); bind_embedding(m); } diff --git a/src/infinicore/pybind11/ops/log_softmax.hpp b/src/infinicore/pybind11/ops/log_softmax.hpp new file mode 100644 index 000000000..3c45bcc1b --- /dev/null +++ b/src/infinicore/pybind11/ops/log_softmax.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include +#include "infinicore/ops/log_softmax.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_log_softmax(py::module &m) { + // 1. 绑定 functional 接口: output = log_softmax(input, dim) + m.def("log_softmax", + &op::log_softmax, + py::arg("input"), + py::arg("dim"), + R"doc(Applies a softmax followed by a logarithm. + + Args: + input (Tensor): The input tensor. + dim (int): A dimension along which log_softmax will be computed. + )doc"); + + // 2. 绑定 explicit output 接口: log_softmax_(output, input, dim) + m.def("log_softmax_", + &op::log_softmax_, + py::arg("output"), + py::arg("input"), + py::arg("dim"), + R"doc(Explicit output LogSoftmax operation. Writes results into output tensor.)doc"); +} + +} // namespace infinicore::ops \ No newline at end of file diff --git a/src/infinicore/pybind11/ops/logaddexp.hpp b/src/infinicore/pybind11/ops/logaddexp.hpp new file mode 100644 index 000000000..08715b368 --- /dev/null +++ b/src/infinicore/pybind11/ops/logaddexp.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include "infinicore/ops/logaddexp.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logaddexp(py::module &m) { + m.def("logaddexp", + &op::logaddexp, + py::arg("a"), + py::arg("b"), + R"doc(Logarithm of the sum of exponentiations of the inputs.)doc"); + m.def("logaddexp_", + &op::logaddexp_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place logaddexp operation. Writes results into c tensor.)doc"); +} + +} // namespace infinicore::ops \ No newline at end of file diff --git a/src/infinicore/pybind11/ops/logaddexp2.hpp b/src/infinicore/pybind11/ops/logaddexp2.hpp new file mode 100644 index 000000000..40a35e71e --- /dev/null +++ b/src/infinicore/pybind11/ops/logaddexp2.hpp @@ -0,0 +1,25 @@ +#pragma once + +#include + +#include "infinicore/ops/logaddexp2.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_logaddexp2(py::module &m) { + m.def("logaddexp2", + &op::logaddexp2, + py::arg("a"), + py::arg("b"), + R"doc(Logarithm of the sum of exponentiations of the inputs in base-2.)doc"); + m.def("logaddexp2_", + &op::logaddexp2_, + py::arg("c"), + py::arg("a"), + py::arg("b"), + R"doc(In-place logaddexp2 operation. Writes results into c tensor.)doc"); +} + +} // namespace infinicore::ops \ No newline at end of file diff --git a/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp b/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp new file mode 100644 index 000000000..167d4c4e3 --- /dev/null +++ b/src/infinicore/pybind11/ops/triplet_margin_with_distance_loss.hpp @@ -0,0 +1,41 @@ +#pragma once + +#include +#include "infinicore/ops/triplet_margin_with_distance_loss.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_triplet_margin_with_distance_loss(py::module &m) { + m.def("triplet_margin_with_distance_loss", + &op::triplet_margin_with_distance_loss, + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin") = 1.0, + py::arg("swap") = false, + py::arg("reduction") = 1, + R"doc(Computes the triplet margin loss with distance. + + Args: + anchor (Tensor): The anchor input tensor. + positive (Tensor): The positive input tensor. + negative (Tensor): The negative input tensor. + margin (float, optional): Default: 1.0. + swap (bool, optional): The distance swap is described in the paper Learning shallow convolutional feature descriptors with triplet losses. Default: False. + reduction (int, optional): Specifies the reduction to apply to the output: 0 (None), 1 (Mean), 2 (Sum). Default: 1. + )doc"); + m.def("triplet_margin_with_distance_loss_", + &op::triplet_margin_with_distance_loss_, + py::arg("output"), + py::arg("anchor"), + py::arg("positive"), + py::arg("negative"), + py::arg("margin"), + py::arg("swap"), + py::arg("reduction"), + R"doc(Explicit output TripletMarginWithDistanceLoss operation. Writes results into output tensor.)doc"); +} + +} // namespace infinicore::ops \ No newline at end of file diff --git a/src/infinicore/pybind11/ops/upsample_nearest.hpp b/src/infinicore/pybind11/ops/upsample_nearest.hpp new file mode 100644 index 000000000..925fba992 --- /dev/null +++ b/src/infinicore/pybind11/ops/upsample_nearest.hpp @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include "infinicore/ops/upsample_nearest.hpp" + +namespace py = pybind11; + +namespace infinicore::ops { + +inline void bind_upsample_nearest(py::module &m) { + // 1. 绑定 functional 接口: output = upsample_nearest(input, output_size) + m.def("upsample_nearest", + &op::upsample_nearest, + py::arg("input"), + py::arg("output_size"), + R"doc(Upsample the input using nearest neighbor interpolation. + + Args: + input (Tensor): The input tensor. + output_size (List[int]): The output spatial size (e.g. [H_out, W_out]). + )doc"); + + // 2. 绑定 explicit output 接口: upsample_nearest_(output, input) + m.def("upsample_nearest_", + &op::upsample_nearest_, + py::arg("output"), + py::arg("input"), + R"doc(Explicit output UpsampleNearest operation. Writes the result into the output tensor.)doc"); +} + +} // namespace infinicore::ops \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.cc b/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.cc new file mode 100644 index 000000000..82204b33c --- /dev/null +++ b/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.cc @@ -0,0 +1,133 @@ +#include "log_softmax_cpu.h" +#include "../../../devices/cpu/common_cpu.h" +#include +#include +#include +#include +#include +#include + +#include "../../../../utils/custom_types.h" + +namespace op::log_softmax::cpu { + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) { + delete _opaque; + _opaque = nullptr; + } +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int dim) { + + auto handle = reinterpret_cast(handle_); + + auto result = LogSoftmaxInfo::create(output_desc, input_desc, dim); + CHECK_RESULT(result); + + *desc_ptr = new Descriptor( + new Opaque(), + result.take(), + 0, + handle->device, + handle->device_id + ); + + return INFINI_STATUS_SUCCESS; +} + +template +void calculate_cpu_impl( + const LogSoftmaxInfo &info, + void *output, + const void *input) { + + size_t outer_size = info.outer_size(); + size_t dim_size = info.dim_size(); + size_t inner_size = info.inner_size(); + + auto out_ptr = reinterpret_cast(output); + auto in_ptr = reinterpret_cast(input); + + size_t total_tasks = outer_size * inner_size; + + #pragma omp parallel for schedule(static) + for (size_t task_id = 0; task_id < total_tasks; ++task_id) { + // 解算当前任务对应的外部索引和内部索引 + size_t o = task_id / inner_size; + size_t i = task_id % inner_size; + + // 计算基地址偏移 + // Layout: [outer, dim, inner] + // Base Offset = o * (dim_size * inner_size) + i + size_t base_offset = o * dim_size * inner_size + i; + size_t stride = inner_size; + std::vector buffer(dim_size); + float max_val = -std::numeric_limits::infinity(); + + for (size_t d = 0; d < dim_size; ++d) { + T val_t = in_ptr[base_offset + d * stride]; + float val = utils::cast(val_t); // 处理 fp16/bf16 + buffer[d] = val; + if (val > max_val) { + max_val = val; + } + } + + // 计算指数和 (Sum) + // Compute sum(exp(x - max)) + float sum_exp = 0.0f; + for (size_t d = 0; d < dim_size; ++d) { + sum_exp += std::exp(buffer[d] - max_val); + } + + // 计算 LogSumExp + // log(sum(e^(x-M))) + M + float log_sum_exp = std::log(sum_exp) + max_val; + + // 计算最终结果并写入 + // output = x - LogSumExp + for (size_t d = 0; d < dim_size; ++d) { + float res = buffer[d] - log_sum_exp; + out_ptr[base_offset + d * stride] = utils::cast(res); + } + } +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + switch (dtype) { + case INFINI_DTYPE_F32: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_F64: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_F16: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_BF16: + cpu::calculate_cpu_impl(_info, output, input); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::log_softmax::cpu \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.h b/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.h new file mode 100644 index 000000000..9ece47dcf --- /dev/null +++ b/src/infiniop/ops/log_softmax/cpu/log_softmax_cpu.h @@ -0,0 +1,8 @@ +#ifndef __LOG_SOFTMAX_CPU_H__ +#define __LOG_SOFTMAX_CPU_H__ + +#include "../log_softmax.h" + +DESCRIPTOR(cpu) + +#endif // __LOG_SOFTMAX_CPU_H__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/cuda/kernel.cuh b/src/infiniop/ops/log_softmax/cuda/kernel.cuh new file mode 100644 index 000000000..ca47cc885 --- /dev/null +++ b/src/infiniop/ops/log_softmax/cuda/kernel.cuh @@ -0,0 +1,140 @@ +#ifndef __LOG_SOFTMAX_CUDA_CUH__ +#define __LOG_SOFTMAX_CUDA_CUH__ + +#include +#include +#include + + +#include +#include +#include + +namespace op::log_softmax::cuda { + +template +__device__ __forceinline__ float to_float(T val) { + return static_cast(val); +} + +// ================================================================== +// Warp Reduction Helpers +// ================================================================== +template +__device__ __forceinline__ T warp_reduce_max(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val = max(val, __shfl_down_sync(0xffffffff, val, offset)); + } + return val; +} + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +// ================================================================== +// Block Reduction Helpers +// ================================================================== +template +__device__ __forceinline__ T block_reduce_max(T val) { + static __shared__ float shared[32]; // Max 32 warps per block + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_max(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + // 假设 BlockDim.x 不超过 1024 (32 warps) + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -INFINITY; + + if (wid == 0) val = warp_reduce_max(val); + + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} + + +template +__global__ void log_softmax_kernel( + T * __restrict__ output, // [Outer, Dim, Inner] + const T * __restrict__ input, // [Outer, Dim, Inner] + size_t dim_size, + size_t inner_size +) { + // 共享内存用于存储 Block Reduction 的结果广播 + __shared__ float s_max; + __shared__ float s_sum; + + unsigned int tid = threadIdx.x; + unsigned int bid = blockIdx.x; + + // 1. 计算当前 Slice 的基地址 + // GridDim.x = Outer * Inner + size_t outer_idx = bid / inner_size; + size_t inner_idx = bid % inner_size; + + // Layout: [outer, dim, inner] + // Base offset = outer * (dim_size * inner_size) + inner_idx + size_t base_offset = outer_idx * dim_size * inner_size + inner_idx; + size_t stride = inner_size; // 元素在 Dim 维度的跨度 + + float local_max = -INFINITY; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + if (val > local_max) { + local_max = val; + } + } + + // Block Reduction 得到全局 Max + float global_max = block_reduce_max(local_max); + if (tid == 0) s_max = global_max; + __syncthreads(); + global_max = s_max; // 广播 + float local_sum = 0.0f; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + local_sum += expf(val - global_max); + } + + // Block Reduction 得到全局 Sum + float global_sum = block_reduce_sum(local_sum); + if (tid == 0) s_sum = global_sum; + __syncthreads(); + global_sum = s_sum; // 广播 + + // 计算 LogSumExp: log(sum) + max + float log_sum_exp = logf(global_sum) + global_max; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + size_t idx = base_offset + i * stride; + float val = to_float(input[idx]); + output[idx] = static_cast(val - log_sum_exp); + } +} + +} // namespace op::log_softmax::cuda + +#endif // __LOG_SOFTMAX_CUDA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/info.h b/src/infiniop/ops/log_softmax/info.h new file mode 100644 index 000000000..0958abcfb --- /dev/null +++ b/src/infiniop/ops/log_softmax/info.h @@ -0,0 +1,84 @@ +#ifndef __LOG_SOFTMAX_INFO_H__ +#define __LOG_SOFTMAX_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::log_softmax { + +class LogSoftmaxInfo { + LogSoftmaxInfo() = default; + +public: + int _dtype; + int _dim; + + size_t _dim_size; + size_t _outer_size; + size_t _inner_size; + + int dtype() const { return _dtype; } + int dim() const { return _dim; } + size_t dim_size() const { return _dim_size; } + size_t outer_size() const { return _outer_size; } + size_t inner_size() const { return _inner_size; } + + LogSoftmaxInfo(int dtype, int dim, size_t dim_size, size_t outer_size, size_t inner_size) + : _dtype(dtype), _dim(dim), + _dim_size(dim_size), _outer_size(outer_size), _inner_size(inner_size) {} + + static utils::Result create( + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int dim) { + + int ndim = input_desc->ndim(); + + if (dim < 0) { + dim += ndim; + } + if (dim < 0 || dim >= ndim) { + return INFINI_STATUS_BAD_PARAM; + } + + size_t dim_size = input_desc->shape()[dim]; + + size_t outer_size = 1; + for (int i = 0; i < dim; ++i) { + outer_size *= input_desc->shape()[i]; + } + + size_t inner_size = 1; + for (int i = dim + 1; i < ndim; ++i) { + inner_size *= input_desc->shape()[i]; + } + + // Validate Shape: LogSoftmax requires input and output shapes to be identical + if (output_desc->ndim() != input_desc->ndim()) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + for (int i = 0; i < ndim; ++i) { + if (output_desc->shape()[i] != input_desc->shape()[i]) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + + // Validate Dtype + if (output_desc->dtype() != input_desc->dtype()) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return utils::Result(LogSoftmaxInfo{ + input_desc->dtype(), + dim, + dim_size, + outer_size, + inner_size + }); + } +}; + +} // namespace op::log_softmax + +#endif // __LOG_SOFTMAX_INFO_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logsoftmax/logsoftmax.h b/src/infiniop/ops/log_softmax/log_softmax.h similarity index 63% rename from src/infiniop/ops/logsoftmax/logsoftmax.h rename to src/infiniop/ops/log_softmax/log_softmax.h index 8babdeab7..22607a8b9 100644 --- a/src/infiniop/ops/logsoftmax/logsoftmax.h +++ b/src/infiniop/ops/log_softmax/log_softmax.h @@ -1,46 +1,48 @@ -#ifndef LOGSOFTMAX_H -#define LOGSOFTMAX_H +#ifndef __LOG_SOFTMAX_H__ +#define __LOG_SOFTMAX_H__ #include "../../operator.h" -#include "info.h" +#include "info.h" // 引用对应的 LogSoftmaxInfo 定义 -#define DESCRIPTOR(NAMESPACE) \ - \ - namespace op::logsoftmax::NAMESPACE { \ - class Descriptor final : public InfiniopDescriptor { \ - struct Opaque; \ - Opaque *_opaque; \ - LogSoftmaxInfo _info; \ - size_t _workspace_size; \ - \ - Descriptor( \ - Opaque *opaque, \ - LogSoftmaxInfo info, \ - size_t workspace_size, \ - infiniDevice_t device_type, \ - int device_id) \ - : InfiniopDescriptor{device_type, device_id}, \ - _opaque(opaque), \ - _info(info), \ - _workspace_size(workspace_size) {} \ - \ - public: \ - ~Descriptor(); \ - \ - size_t workspaceSize() const { return _workspace_size; } \ - \ - static infiniStatus_t create( \ - infiniopHandle_t handle, \ - Descriptor **desc_ptr, \ - infiniopTensorDescriptor_t y_desc, \ - infiniopTensorDescriptor_t x_desc); \ - \ - infiniStatus_t calculate( \ - void *workspace, size_t workspace_size, \ - void *y, \ - const void *x, \ - void *stream) const; \ - }; \ +// 宏定义：用于生成不同命名空间下的 Descriptor 类 +#define DESCRIPTOR(NAMESPACE) \ + namespace op::log_softmax::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + LogSoftmaxInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + Opaque *opaque, \ + LogSoftmaxInfo info, \ + size_t workspace_size, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _info(info), \ + _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t input_desc, \ + int dim); \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *input, \ + void *stream) const; \ + }; \ } -#endif // LOGSOFTMAX_H +#endif // __LOG_SOFTMAX_H__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/metax/log_softmax_metax.h b/src/infiniop/ops/log_softmax/metax/log_softmax_metax.h new file mode 100644 index 000000000..d58085337 --- /dev/null +++ b/src/infiniop/ops/log_softmax/metax/log_softmax_metax.h @@ -0,0 +1,8 @@ +#ifndef __LOG_SOFTMAX_METAX_H__ +#define __LOG_SOFTMAX_METAX_H__ + +#include "../log_softmax.h" + +DESCRIPTOR(metax) + +#endif // __LOG_SOFTMAX_METAX_H__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/metax/log_softmax_metax.maca b/src/infiniop/ops/log_softmax/metax/log_softmax_metax.maca new file mode 100644 index 000000000..c8e27507d --- /dev/null +++ b/src/infiniop/ops/log_softmax/metax/log_softmax_metax.maca @@ -0,0 +1,242 @@ +#include "log_softmax_metax.h" +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_handle.h" +#include +#include +#include +#include +#include +#include +#include + +namespace op::log_softmax::metax { + + +__device__ __forceinline__ float to_float(float val) { return val; } +__device__ __forceinline__ float to_float(double val) { return static_cast(val); } +__device__ __forceinline__ float to_float(__half val) { return __half2float(val); } +__device__ __forceinline__ float to_float(__maca_bfloat16 val) { return __bfloat162float(val); } + +template +__device__ __forceinline__ T warp_reduce_max(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + T shuffled = __shfl_down_sync(0xffffffff, val, offset); + val = (val > shuffled) ? val : shuffled; + } + return val; +} + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +// ================================================================== +// Block Reduction Helpers +// ================================================================== +template +__device__ __forceinline__ T block_reduce_max(T val) { + static __shared__ float shared[32]; // Max 32 warps per block + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_max(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -INFINITY; + + if (wid == 0) val = warp_reduce_max(val); + + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} + +// ================================================================== +// Kernel: LogSoftmax (Online Softmax / 3-Pass Algorithm) +// ================================================================== +template +__global__ void log_softmax_kernel( + T * __restrict__ output, // [Outer, Dim, Inner] + const T * __restrict__ input, // [Outer, Dim, Inner] + size_t dim_size, + size_t inner_size +) { + // 共享内存用于存储 Block Reduction 的结果广播 + __shared__ float s_max; + __shared__ float s_sum; + + unsigned int tid = threadIdx.x; + unsigned int bid = blockIdx.x; + + // 1. 计算当前 Slice 的基地址 + // GridDim.x = Outer * Inner + size_t outer_idx = bid / inner_size; + size_t inner_idx = bid % inner_size; + + // Layout: [outer, dim, inner] + // Base offset = outer * (dim_size * inner_size) + inner_idx + size_t base_offset = outer_idx * dim_size * inner_size + inner_idx; + size_t stride = inner_size; // 元素在 Dim 维度的跨度 + float local_max = -INFINITY; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + if (val > local_max) { + local_max = val; + } + } + + // Block Reduction 得到全局 Max + float global_max = block_reduce_max(local_max); + if (tid == 0) s_max = global_max; + __syncthreads(); + global_max = s_max; // 广播 + float local_sum = 0.0f; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + local_sum += expf(val - global_max); + } + + // Block Reduction 得到全局 Sum + float global_sum = block_reduce_sum(local_sum); + if (tid == 0) s_sum = global_sum; + __syncthreads(); + global_sum = s_sum; // 广播 + + // 计算 LogSumExp: log(sum) + max + float log_sum_exp = logf(global_sum) + global_max; + + // ============================================================ + // Pass 3: Calculate Final Output + // output = x - LogSumExp + // ============================================================ + for (size_t i = tid; i < dim_size; i += blockDim.x) { + size_t idx = base_offset + i * stride; + float val = to_float(input[idx]); + output[idx] = static_cast(val - log_sum_exp); + } +} + +// ================================================================== +// Host Implementation +// ================================================================== + +struct Descriptor::Opaque {}; + +template +void launch_kernel( + void *output, + const void *input, + const LogSoftmaxInfo& info, + void *stream) { + + // 1. 准备指针 + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + auto mc_stream = reinterpret_cast(stream); + + // 2. 准备形状参数 + size_t dim_size = info.dim_size(); + size_t outer_size = info.outer_size(); + size_t inner_size = info.inner_size(); + + // 3. 计算 Grid/Block + // Grid: 总切片数 (Outer * Inner) + // 每个 Block 处理 1 个 Slice (Dim 维度) + size_t total_slices = outer_size * inner_size; + + // Block: 选择一个合理的 Block Size (例如 256) + unsigned int threads_per_block = 256; + + // 根据 dim_size 调整 block size + if (dim_size < 256) { + threads_per_block = 128; + } + if (dim_size < 128) { + threads_per_block = 64; + } + if (dim_size < 64) { + threads_per_block = 32; + } + + // 4. 启动 Kernel + log_softmax_kernel + <<>>( + out_ptr, + in_ptr, + dim_size, + inner_size + ); +} + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int dim) { + + auto handle = reinterpret_cast(handle_); + + auto info_result = LogSoftmaxInfo::create(output_desc, input_desc, dim); + if (!info_result) return info_result.status(); + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel<__half>(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel<__maca_bfloat16>(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::log_softmax::metax \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/moore/log_softmax_moore.h b/src/infiniop/ops/log_softmax/moore/log_softmax_moore.h new file mode 100644 index 000000000..4addf79e0 --- /dev/null +++ b/src/infiniop/ops/log_softmax/moore/log_softmax_moore.h @@ -0,0 +1,8 @@ +#ifndef __LOG_SOFTMAX_MOORE_API_H__ +#define __LOG_SOFTMAX_MOORE_API_H__ + +#include "../log_softmax.h" + +DESCRIPTOR(moore) + +#endif // __LOG_SOFTMAX_MOORE_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/moore/log_softmax_moore.mu b/src/infiniop/ops/log_softmax/moore/log_softmax_moore.mu new file mode 100644 index 000000000..61a5dc441 --- /dev/null +++ b/src/infiniop/ops/log_softmax/moore/log_softmax_moore.mu @@ -0,0 +1,101 @@ +#include "log_softmax_moore.h" +#include "log_softmax_moore_kernel.h" +#include "../../../devices/moore/moore_handle.h" +#include +#include +#include +#include +#include +namespace op::log_softmax::moore { +template +void launch_kernel( + void *output, + const void *input, + const LogSoftmaxInfo& info, + void *stream) { + + // 1. 准备指针 + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + // MUSA 流类型转换 + auto musa_stream = reinterpret_cast(stream); + + // 2. 准备形状参数 + size_t dim_size = info.dim_size(); + size_t outer_size = info.outer_size(); + size_t inner_size = info.inner_size(); + size_t total_slices = outer_size * inner_size; + unsigned int threads_per_block = 256; + + // 如果 dim_size 很小，可以适当减小 block size，但不要小于 32 (Warp Size) + if (dim_size < 256) { + threads_per_block = 128; + } + if (dim_size < 128) { + threads_per_block = 64; + } + if (dim_size < 64) { + threads_per_block = 32; + } + op::log_softmax::moore::log_softmax_kernel + <<>>( + out_ptr, + in_ptr, + dim_size, + inner_size + ); +} +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int dim) { + + auto info_result = LogSoftmaxInfo::create(output_desc, input_desc, dim); + if (!info_result) return info_result.status(); + + // LogSoftmax 此实现为 Online 算法，不需要额外的 Workspace + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + switch (dtype) { + case INFINI_DTYPE_F16: + // MUSA 使用 half + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel<__mt_bfloat16>(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::log_softmax::moore \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/moore/log_softmax_moore_kernel.h b/src/infiniop/ops/log_softmax/moore/log_softmax_moore_kernel.h new file mode 100644 index 000000000..f3429ab28 --- /dev/null +++ b/src/infiniop/ops/log_softmax/moore/log_softmax_moore_kernel.h @@ -0,0 +1,129 @@ +#ifndef __LOG_SOFTMAX_MOORE_H__ +#define __LOG_SOFTMAX_MOORE_H__ + +#include +#include +#include +#include +#include +#include + +namespace op::log_softmax::moore { +template +__device__ __forceinline__ float to_float(T val) { + return static_cast(val); +} +template +__device__ __forceinline__ T warp_reduce_max(T val) { + // 32-thread warp reduction + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val = max(val, __shfl_down_sync(0xffffffff, val, offset)); + } + return val; +} + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +template +__device__ __forceinline__ T block_reduce_max(T val) { + static __shared__ float shared[32]; // Max 32 warps per block + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_max(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : -INFINITY; + + if (wid == 0) val = warp_reduce_max(val); + + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} +template +__global__ void log_softmax_kernel( + T * __restrict__ output, // [Outer, Dim, Inner] + const T * __restrict__ input, // [Outer, Dim, Inner] + size_t dim_size, + size_t inner_size +) { + // 共享内存用于存储 Block Reduction 的结果广播 + __shared__ float s_max; + __shared__ float s_sum; + + unsigned int tid = threadIdx.x; + unsigned int bid = blockIdx.x; + + // 1. 计算当前 Slice 的基地址 + // GridDim.x = Outer * Inner + size_t outer_idx = bid / inner_size; + size_t inner_idx = bid % inner_size; + size_t base_offset = outer_idx * dim_size * inner_size + inner_idx; + size_t stride = inner_size; // 元素在 Dim 维度的跨度 + float local_max = -INFINITY; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + if (val > local_max) { + local_max = val; + } + } + + // Block Reduction 得到全局 Max + float global_max = block_reduce_max(local_max); + // 线程 0 将结果写入共享内存 + if (tid == 0) s_max = global_max; + __syncthreads(); + // 广播到所有线程 + global_max = s_max; + + // ============================================================ + // Pass 2: Calculate Sum of Exponentials + // sum(exp(x - max)) + // ============================================================ + float local_sum = 0.0f; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + float val = to_float(input[base_offset + i * stride]); + local_sum += expf(val - global_max); + } + + // Block Reduction 得到全局 Sum + float global_sum = block_reduce_sum(local_sum); + if (tid == 0) s_sum = global_sum; + __syncthreads(); + global_sum = s_sum; // 广播 + float log_sum_exp = logf(global_sum) + global_max; + for (size_t i = tid; i < dim_size; i += blockDim.x) { + size_t idx = base_offset + i * stride; + float val = to_float(input[idx]); + // 最终写回 + output[idx] = static_cast(val - log_sum_exp); + } +} + +} // namespace op::log_softmax::moore + +#endif // __LOG_SOFTMAX_MOORE_H__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cu b/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cu new file mode 100644 index 000000000..f10fc575c --- /dev/null +++ b/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cu @@ -0,0 +1,115 @@ +#include "log_softmax_nvidia.cuh" +#include "../cuda/kernel.cuh" // 假设这里包含了一些通用的 CUDA 宏或工具 + +#include "../../../handle.h" +#include +#include + +namespace op::log_softmax::nvidia { + +// ================================================================== +// Kernel Launch Logic +// ================================================================== +template +void launch_kernel( + void *output, + const void *input, + const LogSoftmaxInfo& info, + void *stream) { + + // 1. 准备指针 + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + auto cuda_stream = reinterpret_cast(stream); + + // 2. 准备形状参数 + size_t dim_size = info.dim_size(); + size_t outer_size = info.outer_size(); + size_t inner_size = info.inner_size(); + + // 3. 计算 Grid/Block + // Grid: 总切片数 (Outer * Inner) + // 每个 Block 处理 1 个 Slice (Dim 维度) + size_t total_slices = outer_size * inner_size; + + // Block: 选择一个合理的 Block Size (例如 256) + // Kernel 内部使用了循环处理 dim_size > blockDim 的情况， + // 同时使用了 warp reduce，建议 blockDim 至少为 32。 + unsigned int threads_per_block = 256; + + // 如果 dim_size 很小，可以适当减小 block size，但不要小于 32 (Warp Size) + if (dim_size < 256) { + threads_per_block = 128; + } + if (dim_size < 128) { + threads_per_block = 64; + } + if (dim_size < 64) { + threads_per_block = 32; + } + + // 4. 启动 Kernel + // Shared memory 在 kernel 内部静态分配，此处不需要动态分配 + op::log_softmax::cuda::log_softmax_kernel + <<>>( + out_ptr, + in_ptr, + dim_size, + inner_size + ); +} + +// ================================================================== +// Descriptor 实现 +// ================================================================== +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc, + int dim) { + + auto info_result = LogSoftmaxInfo::create(output_desc, input_desc, dim); + if (!info_result) return info_result.status(); + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::log_softmax::nvidia \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cuh b/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cuh new file mode 100644 index 000000000..9a0246e61 --- /dev/null +++ b/src/infiniop/ops/log_softmax/nvidia/log_softmax_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __LOG_SOFTMAX_NVIDIA_CUH__ +#define __LOG_SOFTMAX_NVIDIA_CUH__ + +#include "../log_softmax.h" + +DESCRIPTOR(nvidia) + +#endif // __LOG_SOFTMAX_NVIDIA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/log_softmax/operator.cc b/src/infiniop/ops/log_softmax/operator.cc new file mode 100644 index 000000000..c5039890d --- /dev/null +++ b/src/infiniop/ops/log_softmax/operator.cc @@ -0,0 +1,178 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/log_softmax.h" + +// --- 后端实现头文件 --- +#ifdef ENABLE_CPU_API +#include "cpu/log_softmax_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) +#include "nvidia/log_softmax_nvidia.cuh" +#endif + +#ifdef ENABLE_METAX_API +#include "metax/log_softmax_metax.h" +#endif + +#ifdef ENABLE_MOORE_API +#include "moore/log_softmax_moore.h" +#endif + +extern "C" { + +// ======================================================================= +// 1. 创建算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopCreateLogSoftmaxDescriptor( + infiniopHandle_t handle, + infiniopLogSoftmaxDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input, + int dim) { + + #define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::log_softmax::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output, \ + input, \ + dim) + + switch (handle->device) { + #ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CREATE +} + +// ======================================================================= +// 2. 获取 Workspace 大小 +// ======================================================================= +__C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size) { + + #define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef GET +} + +// ======================================================================= +// 3. 执行计算 (Calculate) +// ======================================================================= +__C infiniStatus_t infiniopLogSoftmax( + infiniopLogSoftmaxDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) { + + #define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, input, stream) + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CALCULATE +} + +// ======================================================================= +// 4. 销毁描述符 +// ======================================================================= +__C infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc) { + + #define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + DELETE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef DELETE +} + +} // extern "C" \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.cc b/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.cc new file mode 100644 index 000000000..9283afa71 --- /dev/null +++ b/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.cc @@ -0,0 +1,43 @@ +#include "logaddexp_cpu.h" + +namespace op::logaddexp::cpu { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate(_info, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp::cpu \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.h b/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.h new file mode 100644 index 000000000..d987639b1 --- /dev/null +++ b/src/infiniop/ops/logaddexp/cpu/logaddexp_cpu.h @@ -0,0 +1,28 @@ +#ifndef __LOGADDEXP_CPU_H__ +#define __LOGADDEXP_CPU_H__ + +#include "../../../elementwise/cpu/elementwise_cpu.h" +#include +#include + +ELEMENTWISE_DESCRIPTOR(logaddexp, cpu) + +namespace op::logaddexp::cpu { + +typedef struct LogAddExpOp { +public: + static constexpr size_t num_inputs = 2; + + template + T operator()(const T &a, const T &b) const { + if (a > b) { + return a + std::log(static_cast(1) + std::exp(b - a)); + } else { + return b + std::log(static_cast(1) + std::exp(a - b)); + } + } +} LogAddExpOp; + +} // namespace op::logaddexp::cpu + +#endif // __LOGADDEXP_CPU_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/cuda/kernel.cuh b/src/infiniop/ops/logaddexp/cuda/kernel.cuh new file mode 100644 index 000000000..7c0807aa8 --- /dev/null +++ b/src/infiniop/ops/logaddexp/cuda/kernel.cuh @@ -0,0 +1,48 @@ +#ifndef __LOGADDEXP_CUDA_H__ +#define __LOGADDEXP_CUDA_H__ + +#include +#include +#include +#include + +namespace op::logaddexp::cuda { +__device__ __forceinline__ float logaddexp_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log1pf(expf(min_val - max_val)); +} +__device__ __forceinline__ double logaddexp_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log1p(exp(min_val - max_val)); +} + +typedef struct LogAddExpOp { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + // half2: 解包为 float2 计算以保证精度 + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp_func(fa.x, fb.x); + res.y = logaddexp_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + // half/bf16: 提升为 float 计算 + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp_func(a, b); + } else { + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } + } +} LogAddExpOp; + +} // namespace op::logaddexp::cuda + +#endif // __LOGADDEXP_CUDA_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/metax/logaddexp_metax.h b/src/infiniop/ops/logaddexp/metax/logaddexp_metax.h new file mode 100644 index 000000000..617bcb98e --- /dev/null +++ b/src/infiniop/ops/logaddexp/metax/logaddexp_metax.h @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP_METAX_API_H__ +#define __LOGADDEXP_METAX_API_H__ + +#include "../../../elementwise/metax/elementwise_metax_api.h" + +ELEMENTWISE_DESCRIPTOR(logaddexp, metax) + +#endif // __LOGADDEXP_METAX_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/metax/logaddexp_metax.maca b/src/infiniop/ops/logaddexp/metax/logaddexp_metax.maca new file mode 100644 index 000000000..2af67056d --- /dev/null +++ b/src/infiniop/ops/logaddexp/metax/logaddexp_metax.maca @@ -0,0 +1,98 @@ +#include "../../../elementwise/metax/elementwise_metax.h" +#include "logaddexp_metax.h" +#include +#include +#include + +namespace op::logaddexp::metax { + +// ================================================================== +// 1. Math Helpers & Functor Definition +// ================================================================== + +__device__ __forceinline__ float logaddexp_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log1pf(expf(min_val - max_val)); +} + +__device__ __forceinline__ double logaddexp_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log1p(exp(min_val - max_val)); +} + +struct LogAddExpOp { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + // half2: 解包为 float2 计算以保证精度 + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp_func(fa.x, fb.x); + res.y = logaddexp_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + // half/bf16: 提升为 float 计算 + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp_func(a, b); + } else { + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } + } +}; + +// ================================================================== +// 2. Descriptor Implementation +// ================================================================== + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + + CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, LogAddExpOp, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, LogAddExpOp, maca_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, LogAddExpOp, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, LogAddExpOp, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::logaddexp::metax \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/moore/logaddexp_moore.h b/src/infiniop/ops/logaddexp/moore/logaddexp_moore.h new file mode 100644 index 000000000..4a12b4ec2 --- /dev/null +++ b/src/infiniop/ops/logaddexp/moore/logaddexp_moore.h @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP_MOORE_API_H__ +#define __LOGADDEXP_MOORE_API_H__ + +#include "../../../elementwise/moore/elementwise_moore_api.h" + +ELEMENTWISE_DESCRIPTOR(logaddexp, moore) + +#endif // __LOGADDEXP_MOORE_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/moore/logaddexp_moore.mu b/src/infiniop/ops/logaddexp/moore/logaddexp_moore.mu new file mode 100644 index 000000000..5fbcdca76 --- /dev/null +++ b/src/infiniop/ops/logaddexp/moore/logaddexp_moore.mu @@ -0,0 +1,48 @@ +#include "../../../elementwise/moore/elementwise_moore.h" +#include "logaddexp_moore.h" +#include "logaddexp_moore_kernel.h" + +namespace op::logaddexp::moore { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, moore::LogAddExpOp, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, moore::LogAddExpOp, __mt_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, moore::LogAddExpOp, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, moore::LogAddExpOp, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp::moore \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/moore/logaddexp_moore_kernel.h b/src/infiniop/ops/logaddexp/moore/logaddexp_moore_kernel.h new file mode 100644 index 000000000..cb5b5ccf2 --- /dev/null +++ b/src/infiniop/ops/logaddexp/moore/logaddexp_moore_kernel.h @@ -0,0 +1,76 @@ +#ifndef __LOGADDEXP_MOORE_KERNEL_H__ +#define __LOGADDEXP_MOORE_KERNEL_H__ + +#include +#include +#include +#include + +namespace op::logaddexp::moore { + +// ================================================================== +// 1. Math Helpers +// ================================================================== +__device__ __forceinline__ float logaddexp_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log1pf(expf(min_val - max_val)); +} + +__device__ __forceinline__ double logaddexp_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log1p(exp(min_val - max_val)); +} + +// ================================================================== +// 2. Functor Definition +// ================================================================== +typedef struct LogAddExpOp { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + // half2: Unpack to float2 for precision + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp_func(fa.x, fb.x); + res.y = logaddexp_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + // half/bf16: Promote to float + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp_func(a, b); + } else { + return static_cast(logaddexp_func(static_cast(a), static_cast(b))); + } + } +} LogAddExpOp; + +// ================================================================== +// 3. Kernel Definition +// ================================================================== +template +__global__ void logaddexp_kernel( + T *output, + const T *a, + const T *b, + size_t n) { + + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + + LogAddExpOp op; + + for (size_t i = idx; i < n; i += stride) { + output[i] = op(a[i], b[i]); + } +} + +} // namespace op::logaddexp::moore + +#endif // __LOGADDEXP_MOORE_KERNEL_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cu b/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cu new file mode 100644 index 000000000..84f1a8481 --- /dev/null +++ b/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cu @@ -0,0 +1,50 @@ +#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" + +#include "../cuda/kernel.cuh" +#include "logaddexp_nvidia.cuh" + +namespace op::logaddexp::nvidia { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, cuda::LogAddExpOp, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, cuda::LogAddExpOp, cuda_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, cuda::LogAddExpOp, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, cuda::LogAddExpOp, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp::nvidia \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cuh b/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cuh new file mode 100644 index 000000000..755d9b105 --- /dev/null +++ b/src/infiniop/ops/logaddexp/nvidia/logaddexp_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP_NVIDIA_CUH__ +#define __LOGADDEXP_NVIDIA_CUH__ + +#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" + +ELEMENTWISE_DESCRIPTOR(logaddexp, nvidia) + +#endif // __LOGADDEXP_NVIDIA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp/operator.cc b/src/infiniop/ops/logaddexp/operator.cc new file mode 100644 index 000000000..1144c3653 --- /dev/null +++ b/src/infiniop/ops/logaddexp/operator.cc @@ -0,0 +1,177 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/logaddexp.h" + +// --- 后端实现头文件 --- +#ifdef ENABLE_CPU_API +#include "cpu/logaddexp_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) +#include "nvidia/logaddexp_nvidia.cuh" +#endif + +#ifdef ENABLE_METAX_API +#include "metax/logaddexp_metax.h" +#endif +#ifdef ENABLE_MOORE_API +#include "moore/logaddexp_moore.h" +#endif + +extern "C" { + +// ======================================================================= +// 1. 创建算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopCreateLogAddExpDescriptor( + infiniopHandle_t handle, + infiniopLogAddExpDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b) { + + #define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::logaddexp::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + c, \ + {a, b}) + + switch (handle->device) { + #ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CREATE +} + +// ======================================================================= +// 2. 获取 Workspace 大小 +// ======================================================================= +__C infiniStatus_t infiniopGetLogAddExpWorkspaceSize(infiniopLogAddExpDescriptor_t desc, size_t *size) { + + #define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef GET +} + +// ======================================================================= +// 3. 执行计算 (Calculate) +// ======================================================================= +__C infiniStatus_t infiniopLogAddExp( + infiniopLogAddExpDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream) { + + #define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, c, {a, b}, stream) + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CALCULATE +} + +// ======================================================================= +// 4. 销毁描述符 +// ======================================================================= +__C infiniStatus_t infiniopDestroyLogAddExpDescriptor(infiniopLogAddExpDescriptor_t desc) { + + #define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + DELETE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef DELETE +} + +} // extern "C" \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.cc b/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.cc new file mode 100644 index 000000000..db1cbf36f --- /dev/null +++ b/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.cc @@ -0,0 +1,47 @@ +#include "logaddexp2_cpu.h" + +namespace op::logaddexp2::cpu { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + + // LogAddExp2 仅支持浮点类型 + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16); + + // create CPU elementwise descriptor + CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate(_info, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate(_info, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp2::cpu \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.h b/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.h new file mode 100644 index 000000000..8383f0f1b --- /dev/null +++ b/src/infiniop/ops/logaddexp2/cpu/logaddexp2_cpu.h @@ -0,0 +1,28 @@ +#ifndef __LOGADDEXP2_CPU_H__ +#define __LOGADDEXP2_CPU_H__ + +#include "../../../elementwise/cpu/elementwise_cpu.h" +#include +#include + +ELEMENTWISE_DESCRIPTOR(logaddexp2, cpu) + +namespace op::logaddexp2::cpu { + +typedef struct LogAddExp2Op { +public: + static constexpr size_t num_inputs = 2; + + template + T operator()(const T &a, const T &b) const { + if (a > b) { + return a + std::log2(static_cast(1) + std::exp2(b - a)); + } else { + return b + std::log2(static_cast(1) + std::exp2(a - b)); + } + } +} LogAddExp2Op; + +} // namespace op::logaddexp2::cpu + +#endif // __LOGADDEXP2_CPU_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/cuda/kernel.cuh b/src/infiniop/ops/logaddexp2/cuda/kernel.cuh new file mode 100644 index 000000000..796f5649b --- /dev/null +++ b/src/infiniop/ops/logaddexp2/cuda/kernel.cuh @@ -0,0 +1,48 @@ +#ifndef __LOGADDEXP2_CUDA_H__ +#define __LOGADDEXP2_CUDA_H__ + +#include +#include +#include +#include + +namespace op::logaddexp2::cuda { + +__device__ __forceinline__ float logaddexp2_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log2f(1.0f + exp2f(min_val - max_val)); +} + +__device__ __forceinline__ double logaddexp2_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log2(1.0 + exp2(min_val - max_val)); +} + +typedef struct LogAddExp2Op { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp2_func(fa.x, fb.x); + res.y = logaddexp2_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp2_func(a, b); + } else { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } + } +} LogAddExp2Op; + +} // namespace op::logaddexp2::cuda + +#endif // __LOGADDEXP2_CUDA_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.h b/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.h new file mode 100644 index 000000000..2e8cec0ce --- /dev/null +++ b/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.h @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP2_METAX_API_H__ +#define __LOGADDEXP2_METAX_API_H__ + +#include "../../../elementwise/metax/elementwise_metax_api.h" + +ELEMENTWISE_DESCRIPTOR(logaddexp2, metax) + +#endif // __LOGADDEXP2_METAX_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.maca b/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.maca new file mode 100644 index 000000000..d7b5c6b59 --- /dev/null +++ b/src/infiniop/ops/logaddexp2/metax/logaddexp2_metax.maca @@ -0,0 +1,96 @@ +#include "../../../elementwise/metax/elementwise_metax.h" +#include "logaddexp2_metax.h" +#include +#include +#include + +namespace op::logaddexp2::metax { + +// ================================================================== +// 1. Math Helpers & Functor Definition +// ================================================================== + +__device__ __forceinline__ float logaddexp2_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log2f(1.0f + exp2f(min_val - max_val)); +} + +__device__ __forceinline__ double logaddexp2_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log2(1.0 + exp2(min_val - max_val)); +} + +struct LogAddExp2Op { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp2_func(fa.x, fb.x); + res.y = logaddexp2_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp2_func(a, b); + } else { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } + } +}; + +// ================================================================== +// 2. Descriptor Implementation +// ================================================================== + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + + CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, LogAddExp2Op, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, LogAddExp2Op, maca_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, LogAddExp2Op, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, LogAddExp2Op, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::logaddexp2::metax \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.h b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.h new file mode 100644 index 000000000..d6bb9a165 --- /dev/null +++ b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.h @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP2_MOORE_API_H__ +#define __LOGADDEXP2_MOORE_API_H__ + +#include "../../../elementwise/moore/elementwise_moore_api.h" + +ELEMENTWISE_DESCRIPTOR(logaddexp2, moore) + +#endif // __LOGADDEXP2_MOORE_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.mu b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.mu new file mode 100644 index 000000000..304ac15af --- /dev/null +++ b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore.mu @@ -0,0 +1,50 @@ +#include "../../../elementwise/moore/elementwise_moore.h" +#include "logaddexp2_moore.h" +#include "logaddexp2_moore_kernel.h" + +namespace op::logaddexp2::moore { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + + CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, moore::LogAddExp2Op, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, moore::LogAddExp2Op, __mt_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, moore::LogAddExp2Op, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, moore::LogAddExp2Op, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp2::moore \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore_kernel.h b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore_kernel.h new file mode 100644 index 000000000..b66276884 --- /dev/null +++ b/src/infiniop/ops/logaddexp2/moore/logaddexp2_moore_kernel.h @@ -0,0 +1,74 @@ +#ifndef __LOGADDEXP2_MOORE_KERNEL_H__ +#define __LOGADDEXP2_MOORE_KERNEL_H__ + +#include +#include +#include +#include + +namespace op::logaddexp2::moore { + +// ================================================================== +// 1. Math Helpers +// ================================================================== +__device__ __forceinline__ float logaddexp2_func(float a, float b) { + float max_val = fmaxf(a, b); + float min_val = fminf(a, b); + return max_val + log2f(1.0f + exp2f(min_val - max_val)); +} + +__device__ __forceinline__ double logaddexp2_func(double a, double b) { + double max_val = fmax(a, b); + double min_val = fmin(a, b); + return max_val + log2(1.0 + exp2(min_val - max_val)); +} + +// ================================================================== +// 2. Functor Definition +// ================================================================== +typedef struct LogAddExp2Op { +public: + static constexpr size_t num_inputs = 2; + + template + __device__ __forceinline__ T operator()(const T &a, const T &b) const { + if constexpr (std::is_same_v) { + float2 fa = __half22float2(a); + float2 fb = __half22float2(b); + float2 res; + res.x = logaddexp2_func(fa.x, fb.x); + res.y = logaddexp2_func(fa.y, fb.y); + return __float22half2_rn(res); + } else if constexpr (std::is_same_v || std::is_same_v) { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } else if constexpr (std::is_same_v) { + return logaddexp2_func(a, b); + } else { + return static_cast(logaddexp2_func(static_cast(a), static_cast(b))); + } + } +} LogAddExp2Op; + +// ================================================================== +// 3. Kernel Definition +// ================================================================== +template +__global__ void logaddexp2_kernel( + T *output, + const T *a, + const T *b, + size_t n) { + + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t stride = blockDim.x * gridDim.x; + + LogAddExp2Op op; + + for (size_t i = idx; i < n; i += stride) { + output[i] = op(a[i], b[i]); + } +} + +} // namespace op::logaddexp2::moore + +#endif // __LOGADDEXP2_MOORE_KERNEL_H__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cu b/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cu new file mode 100644 index 000000000..a3f8ffd0b --- /dev/null +++ b/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cu @@ -0,0 +1,52 @@ +#include "../../../elementwise/nvidia/elementwise_nvidia.cuh" + +#include "../cuda/kernel.cuh" +#include "logaddexp2_nvidia.cuh" + +namespace op::logaddexp2::nvidia { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + std::vector input_desc_vec) { + + auto handle = reinterpret_cast(handle_); + auto dtype = out_desc->dtype(); + CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16, INFINI_DTYPE_F64); + + // create CUDA elementwise descriptor + CREATE_ELEMENTWISE_CUDA_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec) + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + std::vector inputs, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + switch (_dtype) { + case INFINI_DTYPE_F16: + return _device_info->calculate<256, cuda::LogAddExp2Op, half>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_BF16: + return _device_info->calculate<256, cuda::LogAddExp2Op, cuda_bfloat16>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F32: + return _device_info->calculate<256, cuda::LogAddExp2Op, float>(_info, workspace, output, inputs, stream); + case INFINI_DTYPE_F64: + return _device_info->calculate<256, cuda::LogAddExp2Op, double>(_info, workspace, output, inputs, stream); + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} +} // namespace op::logaddexp2::nvidia \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cuh b/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cuh new file mode 100644 index 000000000..1f071dca5 --- /dev/null +++ b/src/infiniop/ops/logaddexp2/nvidia/logaddexp2_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __LOGADDEXP2_NVIDIA_CUH__ +#define __LOGADDEXP2_NVIDIA_CUH__ + +#include "../../../elementwise/nvidia/elementwise_nvidia_api.cuh" + +ELEMENTWISE_DESCRIPTOR(logaddexp2, nvidia) + +#endif // __LOGADDEXP2_NVIDIA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/logaddexp2/operator.cc b/src/infiniop/ops/logaddexp2/operator.cc new file mode 100644 index 000000000..c36fd2410 --- /dev/null +++ b/src/infiniop/ops/logaddexp2/operator.cc @@ -0,0 +1,177 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/logaddexp2.h" + +// --- 后端实现头文件 --- +#ifdef ENABLE_CPU_API +#include "cpu/logaddexp2_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) +#include "nvidia/logaddexp2_nvidia.cuh" +#endif + +#ifdef ENABLE_METAX_API +#include "metax/logaddexp2_metax.h" +#endif +#ifdef ENABLE_MOORE_API +#include "moore/logaddexp2_moore.h" +#endif + +extern "C" { + +// ======================================================================= +// 1. 创建算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopCreateLogAddExp2Descriptor( + infiniopHandle_t handle, + infiniopLogAddExp2Descriptor_t *desc_ptr, + infiniopTensorDescriptor_t c, + infiniopTensorDescriptor_t a, + infiniopTensorDescriptor_t b) { + + #define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::logaddexp2::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + c, \ + {a, b}) + + switch (handle->device) { + #ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CREATE +} + +// ======================================================================= +// 2. 获取 Workspace 大小 +// ======================================================================= +__C infiniStatus_t infiniopGetLogAddExp2WorkspaceSize(infiniopLogAddExp2Descriptor_t desc, size_t *size) { + + #define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef GET +} + +// ======================================================================= +// 3. 执行计算 (Calculate) +// ======================================================================= +__C infiniStatus_t infiniopLogAddExp2( + infiniopLogAddExp2Descriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + void *stream) { + + #define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, c, {a, b}, stream) + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CALCULATE +} + +// ======================================================================= +// 4. 销毁描述符 +// ======================================================================= +__C infiniStatus_t infiniopDestroyLogAddExp2Descriptor(infiniopLogAddExp2Descriptor_t desc) { + + #define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + DELETE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef DELETE +} + +} // extern "C" \ No newline at end of file diff --git a/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.cc b/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.cc deleted file mode 100644 index a6a3876f9..000000000 --- a/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.cc +++ /dev/null @@ -1,130 +0,0 @@ -#include "logsoftmax_cpu.h" -#include "../../../devices/cpu/common_cpu.h" -#include "../../../reduce/cpu/reduce.h" -#include -#include - -namespace op::logsoftmax::cpu { - -Descriptor::~Descriptor() {} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto result = LogSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(result); - *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t logsoftmax(const LogSoftmaxInfo *info, Ty *y, const Tx *x) { -#pragma omp parallel for - for (ptrdiff_t batch = 0; batch < ptrdiff_t(info->batch_size); batch++) { - ptrdiff_t y_offset, x_offset; - - if (info->ndim == 3) { - // For 3D tensors, convert linear batch index back to 2D indices - ptrdiff_t batch_idx = batch / info->seq_len; - ptrdiff_t seq_idx = batch % info->seq_len; - y_offset = batch_idx * info->y_stride_0 + seq_idx * info->y_stride_1; - x_offset = batch_idx * info->x_stride_0 + seq_idx * info->x_stride_1; - } else { - // For 2D tensors, use the flattened strides - y_offset = batch * info->y_stride_b; - x_offset = batch * info->x_stride_b; - } - - Ty *y_ = y + y_offset; - const Tx *x_ = x + x_offset; - - // Find max value for numerical stability - float max_val; - if constexpr (std::is_same::value || std::is_same::value) { - max_val = op::common_cpu::reduce_op::max(x_, info->probs_size, info->x_stride_p); - } else { - max_val = op::common_cpu::reduce_op::max(x_, info->probs_size, info->x_stride_p); - } - - // Compute exp(x - max) and sum - float sum = 0.0f; - for (size_t i = 0; i < info->probs_size; i++) { - float x_val; - if constexpr (std::is_same::value || std::is_same::value) { - x_val = utils::cast(x_[i * info->x_stride_p]); - } else { - x_val = x_[i * info->x_stride_p]; - } - sum += std::exp(x_val - max_val); - } - - // Compute log(sum) - float log_sum = std::log(sum); - - // Compute log_softmax = x - max - log(sum) - for (size_t i = 0; i < info->probs_size; i++) { - float x_val; - if constexpr (std::is_same::value || std::is_same::value) { - x_val = utils::cast(x_[i * info->x_stride_p]); - } else { - x_val = x_[i * info->x_stride_p]; - } - - float result = x_val - max_val - log_sum; - - if constexpr (std::is_same::value || std::is_same::value) { - y_[i * info->y_stride_p] = utils::cast(result); - } else { - y_[i * info->y_stride_p] = result; - } - } - } - - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate( - void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream) const { - - // Handle different input/output dtype combinations - if (_info.x_dtype == INFINI_DTYPE_F16) { - if (_info.y_dtype == INFINI_DTYPE_F16) { - return logsoftmax(&_info, (fp16_t *)y, (const fp16_t *)x); - } else if (_info.y_dtype == INFINI_DTYPE_BF16) { - return logsoftmax(&_info, (bf16_t *)y, (const fp16_t *)x); - } else if (_info.y_dtype == INFINI_DTYPE_F32) { - return logsoftmax(&_info, (float *)y, (const fp16_t *)x); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_info.x_dtype == INFINI_DTYPE_BF16) { - if (_info.y_dtype == INFINI_DTYPE_F16) { - return logsoftmax(&_info, (fp16_t *)y, (const bf16_t *)x); - } else if (_info.y_dtype == INFINI_DTYPE_BF16) { - return logsoftmax(&_info, (bf16_t *)y, (const bf16_t *)x); - } else if (_info.y_dtype == INFINI_DTYPE_F32) { - return logsoftmax(&_info, (float *)y, (const bf16_t *)x); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else if (_info.x_dtype == INFINI_DTYPE_F32) { - if (_info.y_dtype == INFINI_DTYPE_F16) { - return logsoftmax(&_info, (fp16_t *)y, (const float *)x); - } else if (_info.y_dtype == INFINI_DTYPE_BF16) { - return logsoftmax(&_info, (bf16_t *)y, (const float *)x); - } else if (_info.y_dtype == INFINI_DTYPE_F32) { - return logsoftmax(&_info, (float *)y, (const float *)x); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } -} - -} // namespace op::logsoftmax::cpu diff --git a/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.h b/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.h deleted file mode 100644 index 371917bad..000000000 --- a/src/infiniop/ops/logsoftmax/cpu/logsoftmax_cpu.h +++ /dev/null @@ -1,7 +0,0 @@ -#ifndef __LOGSOFTMAX_CPU_H__ -#define __LOGSOFTMAX_CPU_H__ -#include "../logsoftmax.h" - -DESCRIPTOR(cpu) - -#endif diff --git a/src/infiniop/ops/logsoftmax/cuda/kernel.cuh b/src/infiniop/ops/logsoftmax/cuda/kernel.cuh deleted file mode 100644 index d45975075..000000000 --- a/src/infiniop/ops/logsoftmax/cuda/kernel.cuh +++ /dev/null @@ -1,112 +0,0 @@ -#ifndef __LOGSOFTMAX_KERNEL_CUH__ -#define __LOGSOFTMAX_KERNEL_CUH__ - -#include -#include - -template -__device__ void logSoftmaxKernel( - Tdata_out *y, const Tdata_in *x, - size_t batch_size, size_t probs_size, size_t ndim, size_t seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_p, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_p, - ptrdiff_t y_stride_0, ptrdiff_t y_stride_1, - ptrdiff_t x_stride_0, ptrdiff_t x_stride_1) { - - typedef cub::BlockReduce BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - __shared__ Tcompute shared_max_val; - __shared__ Tcompute shared_sum_exp; - - int batch_idx = blockIdx.x; - int tid = threadIdx.x; - - if (batch_idx >= batch_size) { - return; - } - - // Calculate correct memory offsets for 3D tensors - ptrdiff_t y_offset, x_offset; - if (ndim == 3) { - // For 3D tensors, convert linear batch index back to 2D indices - ptrdiff_t batch_dim_idx = batch_idx / seq_len; - ptrdiff_t seq_dim_idx = batch_idx % seq_len; - y_offset = batch_dim_idx * y_stride_0 + seq_dim_idx * y_stride_1; - x_offset = batch_dim_idx * x_stride_0 + seq_dim_idx * x_stride_1; - } else { - // For 2D tensors, use the flattened strides - y_offset = batch_idx * y_stride_b; - x_offset = batch_idx * x_stride_b; - } - - const Tdata_in *x_batch = x + x_offset; - Tdata_out *y_batch = y + y_offset; - - // Find maximum value for numerical stability - Tcompute max_val = static_cast(-INFINITY); - for (int i = tid; i < probs_size; i += BLOCK_SIZE) { - if (i < probs_size) { // Add boundary check - Tcompute val = static_cast(x_batch[i * x_stride_p]); - if constexpr (std::is_same_v) { - max_val = fmaxf(max_val, val); - } else { - max_val = fmax(max_val, val); - } - } - } -#if CUDART_VERSION >= 12090 - max_val = BlockReduce(temp_storage).Reduce(max_val, ::cuda::maximum()); -#else - max_val = BlockReduce(temp_storage).Reduce(max_val, cub::Max()); -#endif - if (tid == 0) { - shared_max_val = max_val; - } - __syncthreads(); - - // Compute sum of exp(x - max) - Tcompute sum_exp = static_cast(0.0); - for (int i = tid; i < probs_size; i += BLOCK_SIZE) { - if (i < probs_size) { // Add boundary check - Tcompute val = static_cast(x_batch[i * x_stride_p]); - if constexpr (std::is_same_v) { - sum_exp += expf(val - shared_max_val); - } else { - sum_exp += exp(val - shared_max_val); - } - } - } - sum_exp = BlockReduce(temp_storage).Sum(sum_exp); - if (tid == 0) { - shared_sum_exp = sum_exp; - } - __syncthreads(); - - // Compute log_softmax = x - max - log(sum_exp) - Tcompute log_sum_exp; - if constexpr (std::is_same_v) { - log_sum_exp = logf(shared_sum_exp); - } else { - log_sum_exp = log(shared_sum_exp); - } - for (int i = tid; i < probs_size; i += BLOCK_SIZE) { - if (i < probs_size) { // Add boundary check - Tcompute val = static_cast(x_batch[i * x_stride_p]); - Tcompute result = val - shared_max_val - log_sum_exp; - y_batch[i * y_stride_p] = static_cast(result); - } - } -} - -template -__global__ void logSoftmax( - Tdata_out *y, const Tdata_in *x, - size_t batch_size, size_t probs_size, size_t ndim, size_t seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_p, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_p, - ptrdiff_t y_stride_0, ptrdiff_t y_stride_1, - ptrdiff_t x_stride_0, ptrdiff_t x_stride_1) { - logSoftmaxKernel(y, x, batch_size, probs_size, ndim, seq_len, y_stride_b, y_stride_p, x_stride_b, x_stride_p, y_stride_0, y_stride_1, x_stride_0, x_stride_1); -} - -#endif // __LOGSOFTMAX_KERNEL_CUH__ diff --git a/src/infiniop/ops/logsoftmax/info.h b/src/infiniop/ops/logsoftmax/info.h deleted file mode 100644 index 10ff7815e..000000000 --- a/src/infiniop/ops/logsoftmax/info.h +++ /dev/null @@ -1,117 +0,0 @@ -#ifndef __LOGSOFTMAX_INFO_H__ -#define __LOGSOFTMAX_INFO_H__ - -#include "../../../utils.h" -#include "../../tensor.h" -#include - -namespace op::logsoftmax { - -class LogSoftmaxInfo { - LogSoftmaxInfo() = default; - -public: - infiniDtype_t x_dtype; - infiniDtype_t y_dtype; - size_t batch_size; - size_t probs_size; - - // Original tensor dimensions for 3D support - size_t ndim; - size_t seq_len; // Only used for 3D tensors - - // Flattened strides for CPU iteration - ptrdiff_t y_stride_b; - ptrdiff_t y_stride_p; - ptrdiff_t x_stride_b; - ptrdiff_t x_stride_p; - - // Original 3D strides for correct memory access - ptrdiff_t y_stride_0, y_stride_1, y_stride_2; - ptrdiff_t x_stride_0, x_stride_1, x_stride_2; - - static utils::Result create(infiniopTensorDescriptor_t y_desc, infiniopTensorDescriptor_t x_desc) { - auto x_dtype = x_desc->dtype(); - auto y_dtype = y_desc->dtype(); - - CHECK_DTYPE(x_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - // Check the output data type, and any dtype is allowed to output fp32. - CHECK_DTYPE(y_dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32); - - auto x_shape = x_desc->shape(); - auto y_shape = y_desc->shape(); - CHECK_SAME_SHAPE(x_shape, y_shape); - - auto ndim = x_desc->ndim(); - if (ndim < 2 || ndim > 3) { - CHECK_STATUS(INFINI_STATUS_BAD_TENSOR_SHAPE); - } - - size_t batch_size, probs_size, seq_len = 0; - if (ndim == 2) { - batch_size = x_shape[0]; - probs_size = x_shape[1]; - } else { // ndim == 3 - batch_size = x_shape[0] * x_shape[1]; - probs_size = x_shape[2]; - seq_len = x_shape[1]; - } - - // Store original strides for all dimensions - ptrdiff_t y_stride_0 = 0, y_stride_1 = 0, y_stride_2 = 0; - ptrdiff_t x_stride_0 = 0, x_stride_1 = 0, x_stride_2 = 0; - - if (ndim == 2) { - y_stride_0 = y_desc->stride(0); // First dimension - y_stride_1 = y_desc->stride(1); // Second dimension - x_stride_0 = x_desc->stride(0); - x_stride_1 = x_desc->stride(1); - } else if (ndim == 3) { - y_stride_0 = y_desc->stride(0); // First dimension (batch) - y_stride_1 = y_desc->stride(1); // Second dimension (seq) - y_stride_2 = y_desc->stride(2); // Third dimension (prob) - x_stride_0 = x_desc->stride(0); - x_stride_1 = x_desc->stride(1); - x_stride_2 = x_desc->stride(2); - } - - ptrdiff_t y_stride_b, y_stride_p, x_stride_b, x_stride_p; - if (ndim == 2) { - y_stride_b = y_desc->stride(0); - y_stride_p = y_desc->stride(1); - x_stride_b = x_desc->stride(0); - x_stride_p = x_desc->stride(1); - } else { // ndim == 3 - // For 3D tensors, flat the first two dimensions - // The CPU implementation expects to iterate through batch_size elements - // where each batch contains probs_size elements - // For flattened iteration, we need stride between consecutive sequences - y_stride_b = y_desc->stride(1); // stride between sequences (20*512 -> 512) - y_stride_p = y_desc->stride(2); // stride within probability dimension - x_stride_b = x_desc->stride(1); // stride between sequences - x_stride_p = x_desc->stride(2); // stride within probability dimension - } - - return utils::Result(LogSoftmaxInfo{ - x_dtype, - y_dtype, - batch_size, - probs_size, - ndim, - seq_len, - y_stride_b, - y_stride_p, - x_stride_b, - x_stride_p, - y_stride_0, - y_stride_1, - y_stride_2, - x_stride_0, - x_stride_1, - x_stride_2}); - } -}; - -} // namespace op::logsoftmax - -#endif // __LOGSOFTMAX_INFO_H__ diff --git a/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cu b/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cu deleted file mode 100644 index 1235b2aaf..000000000 --- a/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cu +++ /dev/null @@ -1,131 +0,0 @@ -#include "../../../devices/nvidia/nvidia_common.cuh" -#include "logsoftmax_nvidia.cuh" - -#include "../../../devices/nvidia/nvidia_kernel_common.cuh" -#include - -#include "../cuda/kernel.cuh" - -namespace op::logsoftmax::nvidia { - -struct Descriptor::Opaque { - std::shared_ptr internal; -}; - -Descriptor::~Descriptor() { - delete _opaque; -} - -infiniStatus_t Descriptor::create( - infiniopHandle_t handle, - Descriptor **desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - auto info = LogSoftmaxInfo::create(y_desc, x_desc); - CHECK_RESULT(info); - *desc_ptr = new Descriptor( - new Opaque{reinterpret_cast(handle)->internal()}, - info.take(), 0, handle->device, handle->device_id); - return INFINI_STATUS_SUCCESS; -} - -template -infiniStatus_t launchKernel(void *y, const void *x, infiniDtype_t x_dtype, infiniDtype_t y_dtype, - size_t batch_size, size_t probs_size, size_t ndim, size_t seq_len, - ptrdiff_t y_stride_b, ptrdiff_t y_stride_p, - ptrdiff_t x_stride_b, ptrdiff_t x_stride_p, - ptrdiff_t y_stride_0, ptrdiff_t y_stride_1, - ptrdiff_t x_stride_0, ptrdiff_t x_stride_1, - cudaStream_t stream) { - dim3 grid(uint32_t(batch_size), 1, 1); - - // Handle mixed precision cases - if (x_dtype == INFINI_DTYPE_F16 && y_dtype == INFINI_DTYPE_F32) { - logSoftmax - <<>>((float *)y, (const half *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_F32 && y_dtype == INFINI_DTYPE_F16) { - logSoftmax - <<>>((half *)y, (const float *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_BF16 && y_dtype == INFINI_DTYPE_F32) { - logSoftmax - <<>>((float *)y, (const __nv_bfloat16 *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_F32 && y_dtype == INFINI_DTYPE_BF16) { - logSoftmax - <<>>((__nv_bfloat16 *)y, (const float *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_F16 && y_dtype == INFINI_DTYPE_F16) { - logSoftmax - <<>>((half *)y, (const half *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_BF16 && y_dtype == INFINI_DTYPE_BF16) { - logSoftmax - <<>>((__nv_bfloat16 *)y, (const __nv_bfloat16 *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else if (x_dtype == INFINI_DTYPE_F32 && y_dtype == INFINI_DTYPE_F32) { - logSoftmax - <<>>((float *)y, (const float *)x, - batch_size, probs_size, ndim, seq_len, - y_stride_b, y_stride_p, - x_stride_b, x_stride_p, - y_stride_0, y_stride_1, - x_stride_0, x_stride_1); - } else { - return INFINI_STATUS_BAD_TENSOR_DTYPE; - } - return INFINI_STATUS_SUCCESS; -} - -infiniStatus_t Descriptor::calculate(void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream_) const { - cudaStream_t stream = (cudaStream_t)stream_; - if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) { - CHECK_STATUS(launchKernel( - y, x, _info.x_dtype, _info.y_dtype, _info.batch_size, _info.probs_size, _info.ndim, _info.seq_len, - _info.y_stride_b, _info.y_stride_p, _info.x_stride_b, _info.x_stride_p, - _info.y_stride_0, _info.y_stride_1, _info.x_stride_0, _info.x_stride_1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) { - CHECK_STATUS(launchKernel( - y, x, _info.x_dtype, _info.y_dtype, _info.batch_size, _info.probs_size, _info.ndim, _info.seq_len, - _info.y_stride_b, _info.y_stride_p, _info.x_stride_b, _info.x_stride_p, - _info.y_stride_0, _info.y_stride_1, _info.x_stride_0, _info.x_stride_1, stream)); - } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) { - CHECK_STATUS(launchKernel( - y, x, _info.x_dtype, _info.y_dtype, _info.batch_size, _info.probs_size, _info.ndim, _info.seq_len, - _info.y_stride_b, _info.y_stride_p, _info.x_stride_b, _info.x_stride_p, - _info.y_stride_0, _info.y_stride_1, _info.x_stride_0, _info.x_stride_1, stream)); - } else { - return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED; - } - return INFINI_STATUS_SUCCESS; -} - -} // namespace op::logsoftmax::nvidia diff --git a/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cuh b/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cuh deleted file mode 100644 index 803143ba7..000000000 --- a/src/infiniop/ops/logsoftmax/nvidia/logsoftmax_nvidia.cuh +++ /dev/null @@ -1,8 +0,0 @@ -#ifndef __LOGSOFTMAX_NVIDIA_H__ -#define __LOGSOFTMAX_NVIDIA_H__ - -#include "../logsoftmax.h" - -DESCRIPTOR(nvidia) - -#endif diff --git a/src/infiniop/ops/logsoftmax/operator.cc b/src/infiniop/ops/logsoftmax/operator.cc deleted file mode 100644 index 7175f5020..000000000 --- a/src/infiniop/ops/logsoftmax/operator.cc +++ /dev/null @@ -1,148 +0,0 @@ -#include "../../operator.h" -#include "../../handle.h" -#include "infiniop/ops/logsoftmax.h" - -#ifdef ENABLE_CPU_API -#include "cpu/logsoftmax_cpu.h" -#endif -#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) -#include "nvidia/logsoftmax_nvidia.cuh" -#endif -#ifdef ENABLE_METAX_API -// #include "metax/logsoftmax_metax.h" -#endif -#ifdef ENABLE_ASCEND_API -// #include "ascend/logsoftmax_ascend.h" -#endif - -__C infiniStatus_t infiniopCreateLogSoftmaxDescriptor( - infiniopHandle_t handle, - infiniopLogSoftmaxDescriptor_t *desc_ptr, - infiniopTensorDescriptor_t y_desc, - infiniopTensorDescriptor_t x_desc) { - -#define CREATE(CASE, NAMESPACE) \ - case CASE: \ - return op::logsoftmax::NAMESPACE::Descriptor::create( \ - handle, \ - reinterpret_cast(desc_ptr), \ - y_desc, \ - x_desc); - - switch (handle->device) { -#ifdef ENABLE_CPU_API - CREATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CREATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - // CREATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CREATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - // CREATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - // CREATE(INFINI_DEVICE_ASCEND, ascend) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__C infiniStatus_t infiniopGetLogSoftmaxWorkspaceSize(infiniopLogSoftmaxDescriptor_t desc, size_t *size) { - -#define GET(CASE, NAMESPACE) \ - case CASE: \ - *size = reinterpret_cast(desc)->workspaceSize(); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - GET(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - GET(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - // GET(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - GET(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - // GET(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - // GET(INFINI_DEVICE_ASCEND, ascend) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__C infiniStatus_t infiniopLogSoftmax( - infiniopLogSoftmaxDescriptor_t desc, - void *workspace, size_t workspace_size, - void *y, - const void *x, - void *stream) { - -#define CALCULATE(CASE, NAMESPACE) \ - case CASE: \ - return reinterpret_cast(desc)->calculate( \ - workspace, workspace_size, y, x, stream); - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - CALCULATE(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - CALCULATE(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - // CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - CALCULATE(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - // CALCULATE(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - // CALCULATE(INFINI_DEVICE_ASCEND, ascend) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} - -__C infiniStatus_t infiniopDestroyLogSoftmaxDescriptor(infiniopLogSoftmaxDescriptor_t desc) { - -#define DESTROY(CASE, NAMESPACE) \ - case CASE: \ - delete reinterpret_cast(desc); \ - return INFINI_STATUS_SUCCESS; - - switch (desc->device_type) { -#ifdef ENABLE_CPU_API - DESTROY(INFINI_DEVICE_CPU, cpu) -#endif -#ifdef ENABLE_NVIDIA_API - DESTROY(INFINI_DEVICE_NVIDIA, nvidia) -#endif -#ifdef ENABLE_ILUVATAR_API - // DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); -#endif -#ifdef ENABLE_QY_API - DESTROY(INFINI_DEVICE_QY, nvidia); -#endif -#ifdef ENABLE_METAX_API - // DESTROY(INFINI_DEVICE_METAX, metax) -#endif -#ifdef ENABLE_ASCEND_API - // DESTROY(INFINI_DEVICE_ASCEND, ascend) -#endif - } - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; -} diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.cc b/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.cc new file mode 100644 index 000000000..262b77a2e --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.cc @@ -0,0 +1,167 @@ +#include "triplet_margin_with_distance_loss_cpu.h" +#include "../../../devices/cpu/common_cpu.h" +#include +#include +#include +#include +#include +#include +#include + +#include "../../../../utils/custom_types.h" + +namespace op::triplet_margin_with_distance_loss::cpu { + +struct Descriptor::Opaque { + size_t batch_size; + size_t feature_dim; +}; + +Descriptor::~Descriptor() { + if (_opaque) { + delete _opaque; + _opaque = nullptr; + } +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t anchor_desc, + infiniopTensorDescriptor_t positive_desc, + infiniopTensorDescriptor_t negative_desc, + float margin, + int swap, + int reduction) { + + auto handle = reinterpret_cast(handle_); + + auto result = TripletMarginWithDistanceLossInfo::create( + output_desc, anchor_desc, positive_desc, negative_desc, margin, swap, reduction); + CHECK_RESULT(result); + + // 解析形状信息 + size_t ndim = anchor_desc->ndim(); + size_t feature_dim = (ndim > 0) ? anchor_desc->shape()[ndim - 1] : 1; + size_t total_elements = result->num_elements(); + size_t batch_size = total_elements / feature_dim; + + auto opaque = new Opaque(); + opaque->batch_size = batch_size; + opaque->feature_dim = feature_dim; + + *desc_ptr = new Descriptor( + opaque, + result.take(), + 0, + handle->device, + handle->device_id + ); + + return INFINI_STATUS_SUCCESS; +} + +// 辅助函数：计算两个向量的欧氏距离 +template +inline float compute_pairwise_distance(const T* x, const T* y, size_t len, float eps = 1e-6f) { + float sum_sq = 0.0f; + for (size_t i = 0; i < len; ++i) { + float diff = utils::cast(x[i]) - utils::cast(y[i]); + sum_sq += diff * diff; + } + return std::sqrt(sum_sq + eps); +} + +// FIX: 移除了 Descriptor::Opaque* 参数，改为直接传入 batch_size 和 feature_dim +template +void calculate_cpu_impl( + const TripletMarginWithDistanceLossInfo &info, + size_t batch_size, + size_t feature_dim, + void *output, + const void *anchor, + const void *positive, + const void *negative) { + + auto out_ptr = reinterpret_cast(output); + auto a_ptr = reinterpret_cast(anchor); + auto p_ptr = reinterpret_cast(positive); + auto n_ptr = reinterpret_cast(negative); + + float margin = info.margin(); + bool swap = info.swap(); + int reduction = info.reduction(); // 0:None, 1:Mean, 2:Sum + + float total_loss = 0.0f; + + #pragma omp parallel for schedule(static) reduction(+:total_loss) + for (size_t i = 0; i < batch_size; ++i) { + size_t offset = i * feature_dim; + + const T* curr_a = a_ptr + offset; + const T* curr_p = p_ptr + offset; + const T* curr_n = n_ptr + offset; + + float dist_pos = compute_pairwise_distance(curr_a, curr_p, feature_dim); + float dist_neg = compute_pairwise_distance(curr_a, curr_n, feature_dim); + + if (swap) { + float dist_pn = compute_pairwise_distance(curr_p, curr_n, feature_dim); + if (dist_pn < dist_neg) { + dist_neg = dist_pn; + } + } + + float loss = std::max(dist_pos - dist_neg + margin, 0.0f); + + if (reduction == 0) { + out_ptr[i] = utils::cast(loss); + } else { + total_loss += loss; + } + } + + if (reduction != 0) { + if (reduction == 1) { // Mean + total_loss /= static_cast(batch_size); + } + out_ptr[0] = utils::cast(total_loss); + } +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream) const { + + auto dtype = _info.dtype(); + // 从 _opaque 中获取形状参数 + size_t batch_size = _opaque->batch_size; + size_t feature_dim = _opaque->feature_dim; + + switch (dtype) { + case INFINI_DTYPE_F32: + cpu::calculate_cpu_impl(_info, batch_size, feature_dim, output, anchor, positive, negative); + break; + case INFINI_DTYPE_F64: + cpu::calculate_cpu_impl(_info, batch_size, feature_dim, output, anchor, positive, negative); + break; + case INFINI_DTYPE_F16: + cpu::calculate_cpu_impl(_info, batch_size, feature_dim, output, anchor, positive, negative); + break; + case INFINI_DTYPE_BF16: + cpu::calculate_cpu_impl(_info, batch_size, feature_dim, output, anchor, positive, negative); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::triplet_margin_with_distance_loss::cpu \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.h b/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.h new file mode 100644 index 000000000..0f862df53 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/cpu/triplet_margin_with_distance_loss_cpu.h @@ -0,0 +1,8 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CPU_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CPU_H__ + +#include "../triplet_margin_with_distance_loss.h" + +DESCRIPTOR(cpu) + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CPU_H__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/cuda/kernel.cuh b/src/infiniop/ops/triplet_margin_with_distance_loss/cuda/kernel.cuh new file mode 100644 index 000000000..1c97141ea --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/cuda/kernel.cuh @@ -0,0 +1,143 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CUDA_CUH__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CUDA_CUH__ + +#include +#include +#include + +#include +#include +#include + +namespace op::triplet_margin_with_distance_loss::cuda { + +// ================================================================== +// 类型转换辅助 +// ================================================================== +__device__ __forceinline__ float to_float(float val) { return val; } +__device__ __forceinline__ float to_float(double val) { return static_cast(val); } +__device__ __forceinline__ float to_float(half val) { return __half2float(val); } +#if !defined(ENABLE_METAX_API) +__device__ __forceinline__ float to_float(nv_bfloat16 val) { return __bfloat162float(val); } +#endif + +// ================================================================== +// Block Reduction Helpers +// ================================================================== +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} + +// ================================================================== +// Kernel: Triplet Margin Loss +// ================================================================== +template +__global__ void triplet_margin_loss_kernel( + T * __restrict__ output, // [BatchSize] (仅当 Reduction=None 时使用) + float * __restrict__ reduction_buffer, // [1] FP32 Accumulator (仅当 Reduction!=None 时使用) + const T * __restrict__ anchor, + const T * __restrict__ positive, + const T * __restrict__ negative, + size_t feature_dim, + float margin, + int swap, + int reduction, // 0: None, 1: Mean, 2: Sum + size_t batch_size +) { + size_t batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; + + size_t tid = threadIdx.x; + size_t stride = blockDim.x; + + size_t offset_base = batch_idx * feature_dim; + + float sum_sq_ap = 0.0f; + float sum_sq_an = 0.0f; + float sum_sq_pn = 0.0f; + + for (size_t i = tid; i < feature_dim; i += stride) { + size_t idx = offset_base + i; + float a = to_float(anchor[idx]); + float p = to_float(positive[idx]); + float n = to_float(negative[idx]); + + float diff_ap = a - p; + sum_sq_ap += diff_ap * diff_ap; + + float diff_an = a - n; + sum_sq_an += diff_an * diff_an; + + if (swap) { + float diff_pn = p - n; + sum_sq_pn += diff_pn * diff_pn; + } + } + + float dist_sq_ap = block_reduce_sum(sum_sq_ap); + float dist_sq_an = block_reduce_sum(sum_sq_an); + float dist_sq_pn = 0.0f; + if (swap) { + dist_sq_pn = block_reduce_sum(sum_sq_pn); + } + + if (tid == 0) { + float eps = 1e-6f; + float dist_ap = sqrtf(dist_sq_ap + eps); + float dist_an = sqrtf(dist_sq_an + eps); + + if (swap) { + float dist_pn = sqrtf(dist_sq_pn + eps); + if (dist_pn < dist_an) { + dist_an = dist_pn; + } + } + + float loss = fmaxf(dist_ap - dist_an + margin, 0.0f); + + if (reduction == 0) { // None + output[batch_idx] = static_cast(loss); + } else { // Sum or Mean + atomicAdd(reduction_buffer, loss); + } + } +} +template +__global__ void cast_and_scale_kernel(T *output, const float *reduction_buffer, size_t batch_size, int reduction) { + if (threadIdx.x == 0) { + float val = reduction_buffer[0]; + + // 如果是 Mean 模式，进行除法 + if (reduction == 1) { + val /= static_cast(batch_size); + } + + output[0] = static_cast(val); + } +} + +} // namespace op::triplet_margin_with_distance_loss::cuda + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_CUDA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/info.h b/src/infiniop/ops/triplet_margin_with_distance_loss/info.h new file mode 100644 index 000000000..b0236ab57 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/info.h @@ -0,0 +1,93 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_INFO_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" + +namespace op::triplet_margin_with_distance_loss { + +class TripletMarginWithDistanceLossInfo { + TripletMarginWithDistanceLossInfo() = default; + +public: + int _dtype; + float _margin; + int _swap; + int _reduction; + size_t _num_elements; + + int dtype() const { return _dtype; } + float margin() const { return _margin; } + int swap() const { return _swap; } + int reduction() const { return _reduction; } + size_t num_elements() const { return _num_elements; } + + TripletMarginWithDistanceLossInfo(int dtype, float margin, int swap, int reduction, size_t num_elements) + : _dtype(dtype), _margin(margin), _swap(swap), _reduction(reduction), _num_elements(num_elements) {} + + static utils::Result create( + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t anchor_desc, + infiniopTensorDescriptor_t positive_desc, + infiniopTensorDescriptor_t negative_desc, + float margin, + int swap, + int reduction) { + + // 1. Validate Dtypes + int dtype = anchor_desc->dtype(); + if (positive_desc->dtype() != dtype || negative_desc->dtype() != dtype || output_desc->dtype() != dtype) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + // 2. Validate Input Shapes + // FIX: 使用 size_t 接收 ndim 以避免符号比较警告 + size_t ndim = anchor_desc->ndim(); + if (positive_desc->ndim() != ndim || negative_desc->ndim() != ndim) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + size_t num_elements = 1; + // FIX: 循环变量使用 size_t + for (size_t i = 0; i < ndim; ++i) { + auto dim_size = anchor_desc->shape()[i]; + if (positive_desc->shape()[i] != dim_size || negative_desc->shape()[i] != dim_size) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + num_elements *= dim_size; + } + + // 3. Validate Output Shape based on Reduction + if (reduction == 0) { // None + if (output_desc->ndim() != ndim) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + for (size_t i = 0; i < ndim; ++i) { + if (output_desc->shape()[i] != anchor_desc->shape()[i]) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + } else { // Mean or Sum + size_t output_size = 1; + // FIX: output_desc->ndim() 返回 size_t，循环变量 i 也应为 size_t + for (size_t i = 0; i < output_desc->ndim(); ++i) { + output_size *= output_desc->shape()[i]; + } + if (output_size != 1) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + } + + return utils::Result(TripletMarginWithDistanceLossInfo{ + dtype, + margin, + swap, + reduction, + num_elements + }); + } +}; + +} // namespace op::triplet_margin_with_distance_loss + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_INFO_H__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.h b/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.h new file mode 100644 index 000000000..962984ade --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.h @@ -0,0 +1,8 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_METAX_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_METAX_H__ + +#include "../triplet_margin_with_distance_loss.h" + +DESCRIPTOR(metax) + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_METAX_H__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.maca b/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.maca new file mode 100644 index 000000000..437fa619f --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/metax/triplet_margin_with_distance_loss_metax.maca @@ -0,0 +1,277 @@ +#include "triplet_margin_with_distance_loss_metax.h" +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_handle.h" +#include +#include +#include +#include +#include +#include +#include + +namespace op::triplet_margin_with_distance_loss::metax { + +// ================================================================== +// Device Helpers: 类型转换与归约 +// ================================================================== + +__device__ __forceinline__ float to_float(float val) { return val; } +__device__ __forceinline__ float to_float(double val) { return static_cast(val); } +__device__ __forceinline__ float to_float(__half val) { return __half2float(val); } +__device__ __forceinline__ float to_float(__maca_bfloat16 val) { return __bfloat162float(val); } + +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} + +// ================================================================== +// Kernels +// ================================================================== + +template +__global__ void triplet_margin_loss_kernel( + T * __restrict__ output, // [BatchSize] (仅当 Reduction=None 时使用) + float * __restrict__ reduction_buffer, // [1] FP32 Accumulator (仅当 Reduction!=None 时使用) + const T * __restrict__ anchor, + const T * __restrict__ positive, + const T * __restrict__ negative, + size_t feature_dim, + float margin, + int swap, + int reduction, // 0: None, 1: Mean, 2: Sum + size_t batch_size +) { + size_t batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; + + size_t tid = threadIdx.x; + size_t stride = blockDim.x; + + size_t offset_base = batch_idx * feature_dim; + + float sum_sq_ap = 0.0f; + float sum_sq_an = 0.0f; + float sum_sq_pn = 0.0f; + + for (size_t i = tid; i < feature_dim; i += stride) { + size_t idx = offset_base + i; + float a = to_float(anchor[idx]); + float p = to_float(positive[idx]); + float n = to_float(negative[idx]); + + float diff_ap = a - p; + sum_sq_ap += diff_ap * diff_ap; + + float diff_an = a - n; + sum_sq_an += diff_an * diff_an; + + if (swap) { + float diff_pn = p - n; + sum_sq_pn += diff_pn * diff_pn; + } + } + + float dist_sq_ap = block_reduce_sum(sum_sq_ap); + float dist_sq_an = block_reduce_sum(sum_sq_an); + float dist_sq_pn = 0.0f; + if (swap) { + dist_sq_pn = block_reduce_sum(sum_sq_pn); + } + + if (tid == 0) { + float eps = 1e-6f; + float dist_ap = sqrtf(dist_sq_ap + eps); + float dist_an = sqrtf(dist_sq_an + eps); + + if (swap) { + float dist_pn = sqrtf(dist_sq_pn + eps); + if (dist_pn < dist_an) { + dist_an = dist_pn; + } + } + + float loss = fmaxf(dist_ap - dist_an + margin, 0.0f); + + if (reduction == 0) { // None + output[batch_idx] = static_cast(loss); + } else { // Sum or Mean + atomicAdd(reduction_buffer, loss); + } + } +} + +template +__global__ void cast_and_scale_kernel(T *output, const float *reduction_buffer, size_t batch_size, int reduction) { + if (threadIdx.x == 0) { + float val = reduction_buffer[0]; + + // 如果是 Mean 模式，进行除法 + if (reduction == 1) { + val /= static_cast(batch_size); + } + + output[0] = static_cast(val); + } +} + +// ================================================================== +// Host Implementation +// ================================================================== + +struct Descriptor::Opaque { + size_t batch_size; + size_t feature_dim; +}; + +template +void launch_kernel( + void *output, + void *workspace, // Workspace pointer (float*) + const void *anchor, + const void *positive, + const void *negative, + const TripletMarginWithDistanceLossInfo& info, + size_t batch_size, + size_t feature_dim, + void *stream) { + + auto out_ptr = reinterpret_cast(output); + auto ws_ptr = reinterpret_cast(workspace); // FP32 Workspace + auto anchor_ptr = reinterpret_cast(anchor); + auto pos_ptr = reinterpret_cast(positive); + auto neg_ptr = reinterpret_cast(negative); + + auto mc_stream = reinterpret_cast(stream); + + float margin = info.margin(); + int swap = info.swap(); + int reduction = info.reduction(); // 0:None, 1:Mean, 2:Sum + + size_t grid_size = batch_size; + + unsigned int threads_per_block = 256; + if (feature_dim < 256) threads_per_block = 128; + if (feature_dim < 128) threads_per_block = 64; + if (feature_dim < 64) threads_per_block = 32; + + // 1. 初始化 Accumulator + if (reduction != 0) { + mcMemsetAsync(ws_ptr, 0, sizeof(float), mc_stream); + } + + triplet_margin_loss_kernel + <<>>( + out_ptr, + ws_ptr, // 传递 workspace + anchor_ptr, + pos_ptr, + neg_ptr, + feature_dim, + margin, + swap, + reduction, + batch_size + ); + + // 3. 后处理: Cast & Mean + if (reduction != 0) { + cast_and_scale_kernel + <<<1, 1, 0, mc_stream>>>( + out_ptr, + ws_ptr, + batch_size, + reduction + ); + } +} + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t anchor_desc, + infiniopTensorDescriptor_t positive_desc, + infiniopTensorDescriptor_t negative_desc, + float margin, + int swap, + int reduction) { + + auto handle = reinterpret_cast(handle_); + + auto info_result = TripletMarginWithDistanceLossInfo::create( + output_desc, anchor_desc, positive_desc, negative_desc, margin, swap, reduction); + if (!info_result) return info_result.status(); + + int ndim = anchor_desc->ndim(); + size_t feature_dim = (ndim > 0) ? anchor_desc->shape()[ndim - 1] : 1; + size_t total_elements = info_result->num_elements(); + size_t batch_size = total_elements / feature_dim; + + auto opaque = new Opaque(); + opaque->batch_size = batch_size; + opaque->feature_dim = feature_dim; + size_t workspace_size = (reduction != 0) ? sizeof(float) : 0; + + *desc_ptr = new Descriptor(opaque, info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream) const { + + auto dtype = _info.dtype(); + size_t batch_size = _opaque->batch_size; + size_t feature_dim = _opaque->feature_dim; + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel<__half>(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel<__maca_bfloat16>(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::triplet_margin_with_distance_loss::metax \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.h b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.h new file mode 100644 index 000000000..57ece38c7 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.h @@ -0,0 +1,8 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_API_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_API_H__ + +#include "../triplet_margin_with_distance_loss.h" + +DESCRIPTOR(moore) + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.mu b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.mu new file mode 100644 index 000000000..ee41d96ac --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore.mu @@ -0,0 +1,149 @@ +#include "triplet_margin_with_distance_loss_moore.h" +#include"triplet_margin_with_distance_loss_moore_kernel.h" +#include "../../../handle.h" +#include +#include +#include +#include +#include + +namespace op::triplet_margin_with_distance_loss::moore { + +struct Descriptor::Opaque { + size_t batch_size; + size_t feature_dim; +}; + +template +void launch_kernel( + void *output, + void *workspace, // Workspace pointer (float*) + const void *anchor, + const void *positive, + const void *negative, + const TripletMarginWithDistanceLossInfo& info, + size_t batch_size, + size_t feature_dim, + void *stream) { + + auto out_ptr = reinterpret_cast(output); + auto ws_ptr = reinterpret_cast(workspace); // FP32 Workspace + auto anchor_ptr = reinterpret_cast(anchor); + auto pos_ptr = reinterpret_cast(positive); + auto neg_ptr = reinterpret_cast(negative); + + // MUSA 流转换 + auto musa_stream = reinterpret_cast(stream); + + float margin = info.margin(); + int swap = info.swap(); + int reduction = info.reduction(); // 0:None, 1:Mean, 2:Sum + + size_t grid_size = batch_size; + + unsigned int threads_per_block = 256; + if (feature_dim < 256) threads_per_block = 128; + if (feature_dim < 128) threads_per_block = 64; + if (feature_dim < 64) threads_per_block = 32; + + // 1. 初始化 Accumulator + if (reduction != 0) { + // 将 float workspace 清零，使用 musaMemsetAsync + musaMemsetAsync(ws_ptr, 0, sizeof(float), musa_stream); + } + + // 2. 启动主 Kernel + // 假设 Kernel 定义在 op::triplet_margin_with_distance_loss::moore 命名空间下 + op::triplet_margin_with_distance_loss::moore::triplet_margin_loss_kernel + <<>>( + out_ptr, + ws_ptr, // 传递 workspace + anchor_ptr, + pos_ptr, + neg_ptr, + feature_dim, + margin, + swap, + reduction, + batch_size + ); + + // 3. 后处理: Cast & Mean + if (reduction != 0) { + op::triplet_margin_with_distance_loss::moore::cast_and_scale_kernel + <<<1, 1, 0, musa_stream>>>( + out_ptr, + ws_ptr, + batch_size, + reduction + ); + } +} + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t anchor_desc, + infiniopTensorDescriptor_t positive_desc, + infiniopTensorDescriptor_t negative_desc, + float margin, + int swap, + int reduction) { + + auto info_result = TripletMarginWithDistanceLossInfo::create( + output_desc, anchor_desc, positive_desc, negative_desc, margin, swap, reduction); + if (!info_result) return info_result.status(); + + int ndim = anchor_desc->ndim(); + size_t feature_dim = (ndim > 0) ? anchor_desc->shape()[ndim - 1] : 1; + size_t total_elements = info_result->num_elements(); + size_t batch_size = total_elements / feature_dim; + + auto opaque = new Opaque(); + opaque->batch_size = batch_size; + opaque->feature_dim = feature_dim; + // Reduction 时需要一个 float 的 workspace 来存累加和 + size_t workspace_size = (reduction != 0) ? sizeof(float) : 0; + + *desc_ptr = new Descriptor(opaque, info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream) const { + + auto dtype = _info.dtype(); + size_t batch_size = _opaque->batch_size; + size_t feature_dim = _opaque->feature_dim; + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel<__mt_bfloat16>(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::triplet_margin_with_distance_loss::moore \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore_kernel.h b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore_kernel.h new file mode 100644 index 000000000..f828c59e1 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/moore/triplet_margin_with_distance_loss_moore_kernel.h @@ -0,0 +1,132 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_KERNEL_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_KERNEL_H__ + +#include +#include +#include + +#include +#include +#include + +namespace op::triplet_margin_with_distance_loss::moore { +__device__ __forceinline__ float to_float(float val) { return val; } +__device__ __forceinline__ float to_float(double val) { return static_cast(val); } +__device__ __forceinline__ float to_float(half val) { return __half2float(val); } +__device__ __forceinline__ float to_float( __mt_bfloat16 val) { return __bfloat162float(val); } +template +__device__ __forceinline__ T warp_reduce_sum(T val) { + for (int offset = 32 / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return val; +} + +template +__device__ __forceinline__ T block_reduce_sum(T val) { + static __shared__ float shared[32]; + int lane = threadIdx.x % 32; + int wid = threadIdx.x / 32; + + val = warp_reduce_sum(val); + + if (lane == 0) shared[wid] = val; + __syncthreads(); + + val = (threadIdx.x < blockDim.x / 32) ? shared[lane] : 0.0f; + + if (wid == 0) val = warp_reduce_sum(val); + + return val; +} + +// ================================================================== +// Kernel: Triplet Margin Loss +// ================================================================== +template +__global__ void triplet_margin_loss_kernel( + T * __restrict__ output, // [BatchSize] (仅当 Reduction=None 时使用) + float * __restrict__ reduction_buffer, // [1] FP32 Accumulator (仅当 Reduction!=None 时使用) + const T * __restrict__ anchor, + const T * __restrict__ positive, + const T * __restrict__ negative, + size_t feature_dim, + float margin, + int swap, + int reduction, // 0: None, 1: Mean, 2: Sum + size_t batch_size +) { + size_t batch_idx = blockIdx.x; + if (batch_idx >= batch_size) return; + + size_t tid = threadIdx.x; + size_t stride = blockDim.x; + + size_t offset_base = batch_idx * feature_dim; + + float sum_sq_ap = 0.0f; + float sum_sq_an = 0.0f; + float sum_sq_pn = 0.0f; + + for (size_t i = tid; i < feature_dim; i += stride) { + size_t idx = offset_base + i; + float a = to_float(anchor[idx]); + float p = to_float(positive[idx]); + float n = to_float(negative[idx]); + + float diff_ap = a - p; + sum_sq_ap += diff_ap * diff_ap; + + float diff_an = a - n; + sum_sq_an += diff_an * diff_an; + + if (swap) { + float diff_pn = p - n; + sum_sq_pn += diff_pn * diff_pn; + } + } + + float dist_sq_ap = block_reduce_sum(sum_sq_ap); + float dist_sq_an = block_reduce_sum(sum_sq_an); + float dist_sq_pn = 0.0f; + if (swap) { + dist_sq_pn = block_reduce_sum(sum_sq_pn); + } + + if (tid == 0) { + float eps = 1e-6f; + float dist_ap = sqrtf(dist_sq_ap + eps); + float dist_an = sqrtf(dist_sq_an + eps); + + if (swap) { + float dist_pn = sqrtf(dist_sq_pn + eps); + if (dist_pn < dist_an) { + dist_an = dist_pn; + } + } + + float loss = fmaxf(dist_ap - dist_an + margin, 0.0f); + + if (reduction == 0) { // None + output[batch_idx] = static_cast(loss); + } else { // Sum or Mean + atomicAdd(reduction_buffer, loss); + } + } +} + +template +__global__ void cast_and_scale_kernel(T *output, const float *reduction_buffer, size_t batch_size, int reduction) { + if (threadIdx.x == 0) { + float val = reduction_buffer[0]; + if (reduction == 1) { + val /= static_cast(batch_size); + } + + output[0] = static_cast(val); + } +} + +} // namespace op::triplet_margin_with_distance_loss::moore + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_MOORE_KERNEL_H__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cu b/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cu new file mode 100644 index 000000000..24917d5cd --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cu @@ -0,0 +1,141 @@ +#include "triplet_margin_with_distance_loss_nvidia.cuh" +#include "../cuda/kernel.cuh" +#include "../../../handle.h" +#include +#include + +namespace op::triplet_margin_with_distance_loss::nvidia { + +struct Descriptor::Opaque { + size_t batch_size; + size_t feature_dim; +}; + +template +void launch_kernel( + void *output, + void *workspace, // Workspace pointer (float*) + const void *anchor, + const void *positive, + const void *negative, + const TripletMarginWithDistanceLossInfo& info, + size_t batch_size, + size_t feature_dim, + void *stream) { + + auto out_ptr = reinterpret_cast(output); + auto ws_ptr = reinterpret_cast(workspace); // FP32 Workspace + auto anchor_ptr = reinterpret_cast(anchor); + auto pos_ptr = reinterpret_cast(positive); + auto neg_ptr = reinterpret_cast(negative); + + auto cuda_stream = reinterpret_cast(stream); + + float margin = info.margin(); + int swap = info.swap(); + int reduction = info.reduction(); // 0:None, 1:Mean, 2:Sum + + size_t grid_size = batch_size; + + unsigned int threads_per_block = 256; + if (feature_dim < 256) threads_per_block = 128; + if (feature_dim < 128) threads_per_block = 64; + if (feature_dim < 64) threads_per_block = 32; + + // 1. 初始化 Accumulator + if (reduction != 0) { + cudaMemsetAsync(ws_ptr, 0, sizeof(float), cuda_stream); + } + + op::triplet_margin_with_distance_loss::cuda::triplet_margin_loss_kernel + <<>>( + out_ptr, + ws_ptr, // 传递 workspace + anchor_ptr, + pos_ptr, + neg_ptr, + feature_dim, + margin, + swap, + reduction, + batch_size + ); + + // 3. 后处理: Cast & Mean + if (reduction != 0) { + op::triplet_margin_with_distance_loss::cuda::cast_and_scale_kernel + <<<1, 1, 0, cuda_stream>>>( + out_ptr, + ws_ptr, + batch_size, + reduction + ); + } +} + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t anchor_desc, + infiniopTensorDescriptor_t positive_desc, + infiniopTensorDescriptor_t negative_desc, + float margin, + int swap, + int reduction) { + + auto info_result = TripletMarginWithDistanceLossInfo::create( + output_desc, anchor_desc, positive_desc, negative_desc, margin, swap, reduction); + if (!info_result) return info_result.status(); + + int ndim = anchor_desc->ndim(); + size_t feature_dim = (ndim > 0) ? anchor_desc->shape()[ndim - 1] : 1; + size_t total_elements = info_result->num_elements(); + size_t batch_size = total_elements / feature_dim; + + auto opaque = new Opaque(); + opaque->batch_size = batch_size; + opaque->feature_dim = feature_dim; + size_t workspace_size = (reduction != 0) ? sizeof(float) : 0; + + *desc_ptr = new Descriptor(opaque, info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream) const { + + auto dtype = _info.dtype(); + size_t batch_size = _opaque->batch_size; + size_t feature_dim = _opaque->feature_dim; + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, workspace, anchor, positive, negative, _info, batch_size, feature_dim, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::triplet_margin_with_distance_loss::nvidia \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cuh b/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cuh new file mode 100644 index 000000000..ff9346ab0 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/nvidia/triplet_margin_with_distance_loss_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_NVIDIA_CUH__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_NVIDIA_CUH__ + +#include "../triplet_margin_with_distance_loss.h" + +DESCRIPTOR(nvidia) + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_NVIDIA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc b/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc new file mode 100644 index 000000000..a583e48b9 --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/operator.cc @@ -0,0 +1,191 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/triplet_margin_with_distance_loss.h" + +// --- 后端实现头文件 --- +#ifdef ENABLE_CPU_API +#include "cpu/triplet_margin_with_distance_loss_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) +#include "nvidia/triplet_margin_with_distance_loss_nvidia.cuh" +#endif + +#ifdef ENABLE_METAX_API +#include "metax/triplet_margin_with_distance_loss_metax.h" +#endif + +#ifdef ENABLE_MOORE_API +#include "moore/triplet_margin_with_distance_loss_moore.h" +#endif + +extern "C" { + +// ======================================================================= +// 1. 创建算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopCreateTripletMarginWithDistanceLossDescriptor( + infiniopHandle_t handle, + infiniopTripletMarginWithDistanceLossDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t anchor, + infiniopTensorDescriptor_t positive, + infiniopTensorDescriptor_t negative, + float margin, + int swap, + int reduction) { + + #define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::triplet_margin_with_distance_loss::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output, \ + anchor, \ + positive, \ + negative, \ + margin, \ + swap, \ + reduction) + + switch (handle->device) { + #ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CREATE +} + +// ======================================================================= +// 2. 获取 Workspace 大小 +// ======================================================================= +__C infiniStatus_t infiniopGetTripletMarginWithDistanceLossWorkspaceSize( + infiniopTripletMarginWithDistanceLossDescriptor_t desc, + size_t *size) { + + #define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef GET +} + +// ======================================================================= +// 3. 执行计算 (Calculate) +// ======================================================================= +__C infiniStatus_t infiniopTripletMarginWithDistanceLoss( + infiniopTripletMarginWithDistanceLossDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *anchor, + const void *positive, + const void *negative, + void *stream) { + + #define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, anchor, positive, negative, stream) + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CALCULATE +} + +// ======================================================================= +// 4. 销毁描述符 +// ======================================================================= +__C infiniStatus_t infiniopDestroyTripletMarginWithDistanceLossDescriptor( + infiniopTripletMarginWithDistanceLossDescriptor_t desc) { + + #define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + DELETE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef DELETE +} + +} // extern "C" \ No newline at end of file diff --git a/src/infiniop/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.h b/src/infiniop/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.h new file mode 100644 index 000000000..b59731fde --- /dev/null +++ b/src/infiniop/ops/triplet_margin_with_distance_loss/triplet_margin_with_distance_loss.h @@ -0,0 +1,52 @@ +#ifndef __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_H__ +#define __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_H__ + +#include "../../operator.h" +#include "info.h" +#define DESCRIPTOR(NAMESPACE) \ + namespace op::triplet_margin_with_distance_loss::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + TripletMarginWithDistanceLossInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + Opaque *opaque, \ + TripletMarginWithDistanceLossInfo info, \ + size_t workspace_size, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _info(info), \ + _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t anchor_desc, \ + infiniopTensorDescriptor_t positive_desc, \ + infiniopTensorDescriptor_t negative_desc, \ + float margin, \ + int swap, \ + int reduction); \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *anchor, \ + const void *positive, \ + const void *negative, \ + void *stream) const; \ + }; \ + } + +#endif // __TRIPLET_MARGIN_WITH_DISTANCE_LOSS_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.cc b/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.cc new file mode 100644 index 000000000..9a13e78fe --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.cc @@ -0,0 +1,170 @@ +#include "upsample_nearest_cpu.h" +#include "../../../devices/cpu/common_cpu.h" +#include +#include +#include +#include + +#include "../../../../utils/custom_types.h" + +namespace op::upsample_nearest::cpu { + +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) { + delete _opaque; + _opaque = nullptr; + } +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t output_desc, + infiniopTensorDescriptor_t input_desc) { + + auto handle = reinterpret_cast(handle_); + + // 创建 Info 对象 + auto result = UpsampleNearestInfo::create(output_desc, input_desc); + CHECK_RESULT(result); + + *desc_ptr = new Descriptor( + new Opaque(), + result.take(), + 0, + handle->device, + handle->device_id + ); + + return INFINI_STATUS_SUCCESS; +} + +// 辅助函数：预计算维度的索引 +// Nearest 插值只需要知道输出坐标对应的输入整数坐标 +std::vector pre_compute_indices( + size_t out_size, + size_t in_size) { + + std::vector indices(out_size); + + // 计算缩放因子 + float scale = static_cast(in_size) / out_size; + + for (size_t i = 0; i < out_size; ++i) { + // Nearest 逻辑：通常向下取整 + // src_idx = floor(dst_idx * scale) + int64_t idx = static_cast(std::floor(i * scale)); + + // 防止越界 (虽理论上不应发生，但为了稳健性) + if (idx >= static_cast(in_size)) { + idx = in_size - 1; + } + indices[i] = idx; + } + return indices; +} + +template +void calculate_cpu_impl( + const UpsampleNearestInfo &info, + void *output, + const void *input) { + + // 获取形状信息 + size_t N = info.n(); + size_t C = info.c(); + size_t in_h = info.h_in(); + size_t in_w = info.w_in(); + size_t out_h = info.h_out(); + size_t out_w = info.w_out(); + + auto out_ptr = reinterpret_cast(output); + auto in_ptr = reinterpret_cast(input); + + // 预计算 H 和 W 维度的索引映射 + auto h_indices = pre_compute_indices(out_h, in_h); + auto w_indices = pre_compute_indices(out_w, in_w); + + size_t n_c = N * C; // 合并 Batch 和 Channel 维度进行并行 + + #pragma omp parallel for schedule(static) + for (size_t nc = 0; nc < n_c; ++nc) { + // 当前 channel 的输入输出起始指针 + const T* src_base = in_ptr + nc * in_h * in_w; + T* dst_base = out_ptr + nc * out_h * out_w; + + for (size_t h = 0; h < out_h; ++h) { + // 获取当前输出行对应的输入行索引 + int64_t src_h = h_indices[h]; + // 缓存该行的输入指针 + const T* src_row = src_base + src_h * in_w; + // 缓存该行的输出指针 + T* dst_row = dst_base + h * out_w; + + for (size_t w = 0; w < out_w; ++w) { + // 获取当前输出列对应的输入列索引 + int64_t src_w = w_indices[w]; + + // 直接赋值 + dst_row[w] = src_row[src_w]; + } + } + } +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + switch (dtype) { + case INFINI_DTYPE_F32: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_F64: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_F16: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_BF16: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_U8: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_I8: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_I16: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_U16: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_I32: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_U32: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_I64: + cpu::calculate_cpu_impl(_info, output, input); + break; + case INFINI_DTYPE_U64: + cpu::calculate_cpu_impl(_info, output, input); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::upsample_nearest::cpu \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.h b/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.h new file mode 100644 index 000000000..51ac2334f --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/cpu/upsample_nearest_cpu.h @@ -0,0 +1,8 @@ +#ifndef __UPSAMPLE_NEAREST_CPU_H__ +#define __UPSAMPLE_NEAREST_CPU_H__ + +#include "../upsample_nearest.h" + +DESCRIPTOR(cpu) + +#endif // __UPSAMPLE_NEAREST_CPU_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/cuda/kernel.cuh b/src/infiniop/ops/upsample_nearest/cuda/kernel.cuh new file mode 100644 index 000000000..380c88ab7 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/cuda/kernel.cuh @@ -0,0 +1,56 @@ +#ifndef __UPSAMPLE_NEAREST_CUDA_CUH__ +#define __UPSAMPLE_NEAREST_CUDA_CUH__ + +#include +#include +#include + +#include +#include + +namespace op::upsample_nearest::cuda { +__device__ __forceinline__ int get_nearest_index( + int out_index, + float scale, + int input_size) { + int idx = static_cast(floorf(out_index * scale)); + return min(max(idx, 0), input_size - 1); +} +template +__global__ void upsample_nearest_kernel( + T * __restrict__ output, // [N, C, H_out, W_out] + const T * __restrict__ input, // [N, C, H_in, W_in] + size_t N, + size_t C, + size_t H_in, + size_t W_in, + size_t H_out, + size_t W_out, + float scale_h, // 预计算的缩放比例 (in_size / out_size) + float scale_w) { // 预计算的缩放比例 (in_size / out_size) + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total_elements = N * C * H_out * W_out; + size_t stride = blockDim.x * gridDim.x; + + for (size_t i = idx; i < total_elements; i += stride) { + // 1. 解构索引 (N, C, H_out, W_out) + // Layout: NCHW + size_t w_out_idx = i % W_out; + size_t temp = i / W_out; + size_t h_out_idx = temp % H_out; + temp /= H_out; + size_t c_idx = temp % C; + size_t n_idx = temp / C; + + // 2. 计算源索引 (Source Indices) + int h_in_idx = get_nearest_index(static_cast(h_out_idx), scale_h, static_cast(H_in)); + int w_in_idx = get_nearest_index(static_cast(w_out_idx), scale_w, static_cast(W_in)); + // Input layout: [N, C, H_in, W_in] + size_t in_offset = (n_idx * C + c_idx) * H_in * W_in + h_in_idx * W_in + w_in_idx; + output[i] = input[in_offset]; + } +} + +} // namespace op::upsample_nearest::cuda + +#endif // __UPSAMPLE_NEAREST_CUDA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/info.h b/src/infiniop/ops/upsample_nearest/info.h new file mode 100644 index 000000000..7ba6df0ba --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/info.h @@ -0,0 +1,118 @@ +#ifndef __UPSAMPLE_NEAREST_INFO_H__ +#define __UPSAMPLE_NEAREST_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" +#include + +namespace op::upsample_nearest { + +class UpsampleNearestInfo { + UpsampleNearestInfo() = default; + +public: + int _dtype; + size_t _n; + size_t _c; + size_t _h_in; + size_t _w_in; + size_t _h_out; + size_t _w_out; + + int dtype() const { return _dtype; } + size_t n() const { return _n; } + size_t c() const { return _c; } + size_t h_in() const { return _h_in; } + size_t w_in() const { return _w_in; } + size_t h_out() const { return _h_out; } + size_t w_out() const { return _w_out; } + + UpsampleNearestInfo(int dtype, + size_t n, size_t c, + size_t h_in, size_t w_in, + size_t h_out, size_t w_out) + : _dtype(dtype), + _n(n), _c(c), + _h_in(h_in), _w_in(w_in), + _h_out(h_out), _w_out(w_out) {} + + static utils::Result create( + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc) { + + size_t ndim = input_desc->ndim(); + // 允许 3D (N, C, W) 和 4D (N, C, H, W) + if (ndim < 3 || ndim > 4) { + // 如果为了兼容性，也可以保留 ndim=2 的逻辑，但通常 upsample 至少有 batch/channel + if (ndim != 2 && ndim != 3 && ndim != 4) + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + if (out_desc->ndim() != ndim) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + if (input_desc->dtype() != out_desc->dtype()) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + size_t n = 1; + size_t c = 1; + size_t h_in = 1, w_in = 1; + size_t h_out = 1, w_out = 1; + + if (ndim == 3) { + // Case: [N, C, W] -> Treat as H=1 + n = input_desc->shape()[0]; + c = input_desc->shape()[1]; + w_in = input_desc->shape()[2]; + + // 检查输出维度一致性 + if (out_desc->shape()[0] != n || out_desc->shape()[1] != c) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + w_out = out_desc->shape()[2]; + + // H 固定为 1 + h_in = 1; + h_out = 1; + } else if (ndim == 4) { + // Case: [N, C, H, W] + n = input_desc->shape()[0]; + c = input_desc->shape()[1]; + h_in = input_desc->shape()[2]; + w_in = input_desc->shape()[3]; + + if (out_desc->shape()[0] != n || out_desc->shape()[1] != c) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + h_out = out_desc->shape()[2]; + w_out = out_desc->shape()[3]; + } else { + // Fallback for ndim=2 or others, previous logic + // Assuming [H, W] or similar + for (size_t i = 0; i < ndim - 2; ++i) { + if (input_desc->shape()[i] != out_desc->shape()[i]) return INFINI_STATUS_BAD_TENSOR_SHAPE; + c *= input_desc->shape()[i]; + } + h_in = input_desc->shape()[ndim - 2]; + w_in = input_desc->shape()[ndim - 1]; + h_out = out_desc->shape()[ndim - 2]; + w_out = out_desc->shape()[ndim - 1]; + } + + if (h_in == 0 || w_in == 0 || h_out == 0 || w_out == 0) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + return utils::Result(UpsampleNearestInfo{ + input_desc->dtype(), + n, c, + h_in, w_in, + h_out, w_out + }); + } +}; + +} // namespace op::upsample_nearest + +#endif // __UPSAMPLE_NEAREST_INFO_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.h b/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.h new file mode 100644 index 000000000..882d5d61b --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.h @@ -0,0 +1,8 @@ +#ifndef __UPSAMPLE_NEAREST_METAX_H__ +#define __UPSAMPLE_NEAREST_METAX_H__ + +#include "../upsample_nearest.h" + +DESCRIPTOR(metax) + +#endif // __UPSAMPLE_NEAREST_METAX_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.maca b/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.maca new file mode 100644 index 000000000..f1741c4a0 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/metax/upsample_nearest_metax.maca @@ -0,0 +1,207 @@ +#include "upsample_nearest_metax.h" +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_handle.h" + +#include +#include +#include + +#include +#include +#include +#include + +namespace op::upsample_nearest::metax { + +// ================================================================== +// 1. Device Kernel Implementation +// ================================================================== + +__device__ __forceinline__ int get_nearest_index( + int out_index, + float scale, + int input_size) { + // 使用 floorf 计算最近邻索引 + int idx = static_cast(floorf(out_index * scale)); + // 边界钳制，防止索引越界 + return min(max(idx, 0), input_size - 1); +} + +template +__global__ void upsample_nearest_kernel( + T * __restrict__ output, // [N, C, H_out, W_out] + const T * __restrict__ input, // [N, C, H_in, W_in] + size_t N, + size_t C, + size_t H_in, + size_t W_in, + size_t H_out, + size_t W_out, + float scale_h, // 预计算的缩放比例 (in_size / out_size) + float scale_w) { // 预计算的缩放比例 (in_size / out_size) + + // Grid-Stride Loop: 处理每一个输出元素 + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total_elements = N * C * H_out * W_out; + size_t stride = blockDim.x * gridDim.x; + + for (size_t i = idx; i < total_elements; i += stride) { + // 1. 解构索引 (N, C, H_out, W_out) + // Layout: NCHW + size_t w_out_idx = i % W_out; + size_t temp = i / W_out; + size_t h_out_idx = temp % H_out; + temp /= H_out; + size_t c_idx = temp % C; + size_t n_idx = temp / C; + + // 2. 计算源索引 (Source Indices) + int h_in_idx = get_nearest_index(static_cast(h_out_idx), scale_h, static_cast(H_in)); + int w_in_idx = get_nearest_index(static_cast(w_out_idx), scale_w, static_cast(W_in)); + + // 3. 计算输入数据的线性偏移量 + // Input layout: [N, C, H_in, W_in] + size_t in_offset = (n_idx * C + c_idx) * H_in * W_in + h_in_idx * W_in + w_in_idx; + + // 4. 读取并写入数据 (直接赋值，无插值) + output[i] = input[in_offset]; + } +} + +// ================================================================== +// 2. Host Launch Logic +// ================================================================== + +template +void launch_kernel( + void *output, + const void *input, + const UpsampleNearestInfo& info, + void *stream) { + + // 1. Prepare Pointers + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + // MACA stream conversion + auto mc_stream = reinterpret_cast(stream); + + // 2. Prepare Dimensions + size_t N = info.n(); + size_t C = info.c(); + size_t H_in = info.h_in(); + size_t W_in = info.w_in(); + size_t H_out = info.h_out(); + size_t W_out = info.w_out(); + + // 3. Pre-compute Scaling Factors on Host + // Nearest neighbor scaling: in_size / out_size + float scale_h = static_cast(H_in) / H_out; + float scale_w = static_cast(W_in) / W_out; + + // 4. Configure Grid/Block + // Total number of output elements + size_t total_elements = N * C * H_out * W_out; + size_t block_size = 256; + size_t grid_size = (total_elements + block_size - 1) / block_size; + + // Cap grid size to avoid launch failures on huge tensors + // MetaX/CUDA grid limitation + if (grid_size > 65535) grid_size = 65535; + + upsample_nearest_kernel + <<>>( + out_ptr, + in_ptr, + N, C, H_in, W_in, H_out, W_out, + scale_h, scale_w + ); +} + +// ================================================================== +// 3. Descriptor Implementation +// ================================================================== +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc) { + + auto handle_ptr = reinterpret_cast(handle); + auto info_result = UpsampleNearestInfo::create(out_desc, input_desc); + if (!info_result) return info_result.status(); + + // No extra workspace needed for this op + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle_ptr->device, handle_ptr->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + // Verify pointers + if (!output || !input) { + return INFINI_STATUS_BAD_PARAM; + } + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel<__half>(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + // 使用 MACA 的 bfloat16 类型 + launch_kernel<__maca_bfloat16>(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + // Nearest Neighbor 插值通常也支持整型 + case INFINI_DTYPE_U8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I64: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::upsample_nearest::metax \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.h b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.h new file mode 100644 index 000000000..90d217604 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.h @@ -0,0 +1,8 @@ +#ifndef __UPSAMPLE_NEAREST_MOORE_API_H__ +#define __UPSAMPLE_NEAREST_MOORE_API_H__ + +#include "../upsample_nearest.h" + +DESCRIPTOR(moore) + +#endif // __UPSAMPLE_NEAREST_MOORE_API_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.mu b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.mu new file mode 100644 index 000000000..c53cf7523 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore.mu @@ -0,0 +1,144 @@ +#include "upsample_nearest_moore.h" +#include "upsample_nearest_moore_kernel.h" +#include "../../../handle.h" +#include +#include +#include +#include +#include + +namespace op::upsample_nearest::moore { + +// ================================================================== +// Kernel Launch Logic +// ================================================================== +template +void launch_kernel( + void *output, + const void *input, + const UpsampleNearestInfo& info, + void *stream) { + + // 1. Prepare Pointers + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + auto musa_stream = reinterpret_cast(stream); + + // 2. Prepare Dimensions + size_t N = info.n(); + size_t C = info.c(); + size_t H_in = info.h_in(); + size_t W_in = info.w_in(); + size_t H_out = info.h_out(); + size_t W_out = info.w_out(); + + // 3. Pre-compute Scaling Factors on Host + // Nearest neighbor scaling: in_size / out_size + float scale_h = static_cast(H_in) / H_out; + float scale_w = static_cast(W_in) / W_out; + + // 4. Configure Grid/Block + // Total number of output elements + size_t total_elements = N * C * H_out * W_out; + size_t block_size = 256; + size_t grid_size = (total_elements + block_size - 1) / block_size; + + // Cap grid size to avoid launch failures on huge tensors (handling via grid-stride loop) + if (grid_size > 65535) grid_size = 65535; + + op::upsample_nearest::moore::upsample_nearest_kernel + <<>>( + out_ptr, + in_ptr, + N, C, H_in, W_in, H_out, W_out, + scale_h, scale_w + ); +} + +// ================================================================== +// Descriptor Implementation +// ================================================================== +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc) { + + auto info_result = UpsampleNearestInfo::create(out_desc, input_desc); + if (!info_result) return info_result.status(); + + // No extra workspace needed for this op + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + // Verify pointers + if (!output || !input) { + return INFINI_STATUS_BAD_PARAM; + } + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + // Moore 架构下 BF16 类型 + launch_kernel<__mt_bfloat16>(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + // 整型支持 + case INFINI_DTYPE_U8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I64: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::upsample_nearest::moore \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore_kernel.h b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore_kernel.h new file mode 100644 index 000000000..1923e0d96 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/moore/upsample_nearest_moore_kernel.h @@ -0,0 +1,55 @@ +#ifndef __UPSAMPLE_NEAREST_MOORE_KERNEL_H__ +#define __UPSAMPLE_NEAREST_MOORE_KERNEL_H__ +#include +#include +#include +#include +#include + +namespace op::upsample_nearest::moore { +__device__ __forceinline__ int get_nearest_index( + int out_index, + float scale, + int input_size) { + int idx = static_cast(floorf(out_index * scale)); + return min(max(idx, 0), input_size - 1); +} +template +__global__ void upsample_nearest_kernel( + T * __restrict__ output, // [N, C, H_out, W_out] + const T * __restrict__ input, // [N, C, H_in, W_in] + size_t N, + size_t C, + size_t H_in, + size_t W_in, + size_t H_out, + size_t W_out, + float scale_h, // 预计算的缩放比例 (in_size / out_size) + float scale_w) { // 预计算的缩放比例 (in_size / out_size) + + // Grid-Stride Loop: 处理每一个输出元素 + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + size_t total_elements = N * C * H_out * W_out; + size_t stride = blockDim.x * gridDim.x; + + for (size_t i = idx; i < total_elements; i += stride) { + // 1. 解构索引 (N, C, H_out, W_out) + // Layout: NCHW + size_t w_out_idx = i % W_out; + size_t temp = i / W_out; + size_t h_out_idx = temp % H_out; + temp /= H_out; + size_t c_idx = temp % C; + size_t n_idx = temp / C; + + // 2. 计算源索引 (Source Indices) + int h_in_idx = get_nearest_index(static_cast(h_out_idx), scale_h, static_cast(H_in)); + int w_in_idx = get_nearest_index(static_cast(w_out_idx), scale_w, static_cast(W_in)); + size_t in_offset = (n_idx * C + c_idx) * H_in * W_in + h_in_idx * W_in + w_in_idx; + output[i] = input[in_offset]; + } +} + +} // namespace op::upsample_nearest::moore + +#endif // __UPSAMPLE_NEAREST_MOORE_KERNEL_H__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cu b/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cu new file mode 100644 index 000000000..5e552ebe2 --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cu @@ -0,0 +1,145 @@ +#include "upsample_nearest_nvidia.cuh" +#include "../cuda/kernel.cuh" +#include "../../../handle.h" +#include +#include + +namespace op::upsample_nearest::nvidia { + +template +static inline bool is_aligned(const void *ptr, size_t alignment) { + return reinterpret_cast(ptr) % alignment == 0; +} + +// ================================================================== +// Kernel Launch Logic +// ================================================================== +template +void launch_kernel( + void *output, + const void *input, + const UpsampleNearestInfo& info, + void *stream) { + + // 1. Prepare Pointers + auto in_ptr = reinterpret_cast(input); + auto out_ptr = reinterpret_cast(output); + + auto cuda_stream = reinterpret_cast(stream); + + // 2. Prepare Dimensions + size_t N = info.n(); + size_t C = info.c(); + size_t H_in = info.h_in(); + size_t W_in = info.w_in(); + size_t H_out = info.h_out(); + size_t W_out = info.w_out(); + + // 3. Pre-compute Scaling Factors on Host + // Nearest neighbor scaling: in_size / out_size + float scale_h = static_cast(H_in) / H_out; + float scale_w = static_cast(W_in) / W_out; + + // 4. Configure Grid/Block + // Total number of output elements + size_t total_elements = N * C * H_out * W_out; + size_t block_size = 256; + size_t grid_size = (total_elements + block_size - 1) / block_size; + + // Cap grid size to avoid launch failures on huge tensors + if (grid_size > 65535) grid_size = 65535; + + op::upsample_nearest::cuda::upsample_nearest_kernel + <<>>( + out_ptr, + in_ptr, + N, C, H_in, W_in, H_out, W_out, + scale_h, scale_w + ); +} + +// ================================================================== +// Descriptor Implementation +// ================================================================== +struct Descriptor::Opaque {}; + +Descriptor::~Descriptor() { + if (_opaque) delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t out_desc, + infiniopTensorDescriptor_t input_desc) { + + auto info_result = UpsampleNearestInfo::create(out_desc, input_desc); + if (!info_result) return info_result.status(); + + // No extra workspace needed for this op + size_t workspace_size = 0; + + *desc_ptr = new Descriptor(new Opaque(), info_result.take(), workspace_size, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) const { + + auto dtype = _info.dtype(); + + // Verify pointers + if (!output || !input) { + return INFINI_STATUS_BAD_PARAM; + } + + switch (dtype) { + case INFINI_DTYPE_F16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_BF16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_F64: + launch_kernel(output, input, _info, stream); + break; + // Nearest Neighbor 插值通常也支持整型 (如 Mask 处理) + case INFINI_DTYPE_U8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I8: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U16: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U32: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_I64: + launch_kernel(output, input, _info, stream); + break; + case INFINI_DTYPE_U64: + launch_kernel(output, input, _info, stream); + break; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::upsample_nearest::nvidia \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cuh b/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cuh new file mode 100644 index 000000000..45817fe1c --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/nvidia/upsample_nearest_nvidia.cuh @@ -0,0 +1,7 @@ +#ifndef __UPSAMPLE_NEAREST_NVIDIA_CUH__ +#define __UPSAMPLE_NEAREST_NVIDIA_CUH__ + +#include "../upsample_nearest.h" +DESCRIPTOR(nvidia) + +#endif // __UPSAMPLE_NEAREST_NVIDIA_CUH__ \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/operator.cc b/src/infiniop/ops/upsample_nearest/operator.cc new file mode 100644 index 000000000..99241982e --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/operator.cc @@ -0,0 +1,176 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/upsample_nearest.h" + +// --- 后端实现头文件 --- +#ifdef ENABLE_CPU_API +#include "cpu/upsample_nearest_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) +#include "nvidia/upsample_nearest_nvidia.cuh" +#endif + +#ifdef ENABLE_METAX_API +#include "metax/upsample_nearest_metax.h" +#endif + +#ifdef ENABLE_MOORE_API +#include "moore/upsample_nearest_moore.h" +#endif + +extern "C" { + +// ======================================================================= +// 1. 创建算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopCreateUpsampleNearestDescriptor( + infiniopHandle_t handle, + infiniopUpsampleNearestDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t output, + infiniopTensorDescriptor_t input) { + + #define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::upsample_nearest::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + output, \ + input) + + switch (handle->device) { + #ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CREATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CREATE +} + +// ======================================================================= +// 2. 获取 Workspace 大小 +// ======================================================================= +__C infiniStatus_t infiniopGetUpsampleNearestWorkspaceSize(infiniopUpsampleNearestDescriptor_t desc, size_t *size) { + + #define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + GET(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + GET(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef GET +} + +// ======================================================================= +// 3. 执行计算 (Calculate) +// ======================================================================= +__C infiniStatus_t infiniopUpsampleNearest( + infiniopUpsampleNearestDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *output, + const void *input, + void *stream) { + + #define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, output, input, stream) + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + CALCULATE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef CALCULATE +} + +// ======================================================================= +// 4. 销毁算子描述符 +// ======================================================================= +__C infiniStatus_t infiniopDestroyUpsampleNearestDescriptor(infiniopUpsampleNearestDescriptor_t desc) { + + #define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + #ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); + #endif + #ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); + #endif + #ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); + #endif + #ifdef ENABLE_QY_API + DELETE(INFINI_DEVICE_QY, nvidia); + #endif + #ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); + #endif + #ifdef ENABLE_MOORE_API + DELETE(INFINI_DEVICE_MOORE, moore); + #endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + #undef DELETE +} + +} // extern "C" \ No newline at end of file diff --git a/src/infiniop/ops/upsample_nearest/upsample_nearest.h b/src/infiniop/ops/upsample_nearest/upsample_nearest.h new file mode 100644 index 000000000..66f6074eb --- /dev/null +++ b/src/infiniop/ops/upsample_nearest/upsample_nearest.h @@ -0,0 +1,46 @@ +#ifndef __UPSAMPLE_NEAREST_H__ +#define __UPSAMPLE_NEAREST_H__ + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + namespace op::upsample_nearest::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + UpsampleNearestInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + Opaque *opaque, \ + UpsampleNearestInfo info, \ + size_t workspace_size, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _info(info), \ + _workspace_size(workspace_size) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t output_desc, \ + infiniopTensorDescriptor_t input_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, \ + size_t workspace_size, \ + void *output, \ + const void *input, \ + void *stream) const; \ + }; \ + } + +#endif // __UPSAMPLE_NEAREST_H__ \ No newline at end of file diff --git a/test/infinicore/ops/log_softmax.py b/test/infinicore/ops/log_softmax.py index 68af20f1e..eb101670b 100644 --- a/test/infinicore/ops/log_softmax.py +++ b/test/infinicore/ops/log_softmax.py @@ -71,9 +71,8 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.log_softmax(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): - # """InfiniCore implementation (operator not yet available).""" - # return infinicore.nn.functional.log_softmax(*args, **kwargs) + def infinicore_operator(self, *args, **kwargs): + return infinicore.nn.functional.log_softmax(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/logaddexp.py b/test/infinicore/ops/logaddexp.py index 0266e8276..f787b9170 100644 --- a/test/infinicore/ops/logaddexp.py +++ b/test/infinicore/ops/logaddexp.py @@ -102,9 +102,8 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.logaddexp(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): - # """InfiniCore implementation (operator not yet available).""" - # return infinicore.logaddexp(*args, **kwargs) + def infinicore_operator(self, *args, **kwargs): + return infinicore.logaddexp(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/logaddexp2.py b/test/infinicore/ops/logaddexp2.py index 6ee1b984c..9195b54f6 100644 --- a/test/infinicore/ops/logaddexp2.py +++ b/test/infinicore/ops/logaddexp2.py @@ -102,9 +102,8 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.logaddexp2(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): - # """InfiniCore implementation (operator not yet available).""" - # return infinicore.logaddexp2(*args, **kwargs) + def infinicore_operator(self, *args, **kwargs): + return infinicore.logaddexp2(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/triplet_margin_with_distance_loss.py b/test/infinicore/ops/triplet_margin_with_distance_loss.py index 35ca4a1ea..dfd86c798 100644 --- a/test/infinicore/ops/triplet_margin_with_distance_loss.py +++ b/test/infinicore/ops/triplet_margin_with_distance_loss.py @@ -70,9 +70,8 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.triplet_margin_with_distance_loss(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): - # """InfiniCore implementation (operator not yet available).""" - # return infinicore.nn.functional.triplet_margin_with_distance_loss(*args, **kwargs) + def infinicore_operator(self, *args, **kwargs): + return infinicore.nn.functional.triplet_margin_with_distance_loss(*args, **kwargs) def main(): diff --git a/test/infinicore/ops/upsample_nearest.py b/test/infinicore/ops/upsample_nearest.py index 58c6d4e96..045f56218 100644 --- a/test/infinicore/ops/upsample_nearest.py +++ b/test/infinicore/ops/upsample_nearest.py @@ -75,9 +75,8 @@ def get_test_cases(self): def torch_operator(self, *args, **kwargs): return torch.nn.functional.interpolate(*args, **kwargs) - # def infinicore_operator(self, *args, **kwargs): - # """InfiniCore implementation (operator not yet available).""" - # return infinicore.nn.functional.interpolate(*args, **kwargs) + def infinicore_operator(self, *args, **kwargs): + return infinicore.nn.functional.interpolate(*args, **kwargs) def main():