mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-18 18:52:16 +00:00
refactor Cuda Ops Sum, Max, Min, remove dup code (#1946)
refactor Cuda Ops Sum, Max, Min, remove dup code
This commit is contained in:
parent
7b39f5090c
commit
00e24ae4fe
2 changed files with 89 additions and 149 deletions
|
|
@ -223,33 +223,37 @@ BINARY_LOGICALOP_TYPED(Or, 7, bool)
|
|||
BINARY_LOGICALOP_TYPED(Xor, 7, bool)
|
||||
BINARY_OP_HFD(PRelu, 7)
|
||||
|
||||
template <typename T>
|
||||
Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
typedef typename ToCudaType<T>::MappedType CudaT;
|
||||
template <typename T, typename CudaT>
|
||||
Status VariadicInputBase<T, CudaT>::ComputeMethod(OpKernelContext* context, ImplCompute Impl_Compute) const {
|
||||
const auto& node = Node();
|
||||
const auto& node_name = node.Name();
|
||||
auto input_count = node.InputArgCount().front();
|
||||
ORT_RETURN_IF_NOT(input_count >= 1, "Must have 1 or more inputs");
|
||||
|
||||
auto lhs_tensor = context->Input<Tensor>(0);
|
||||
if (input_count == 1) {
|
||||
auto input_tensor = context->Input<Tensor>(0);
|
||||
const auto& input_shape = input_tensor->Shape();
|
||||
const auto& input_shape = lhs_tensor->Shape();
|
||||
auto output_tensor = context->Output(0, input_shape);
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice));
|
||||
if (lhs_tensor->DataRaw() != output_tensor->DataRaw()) {
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), lhs_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice));
|
||||
}
|
||||
} else {
|
||||
// compute output shape first, using broadcast rule
|
||||
TensorShape output_shape;
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input<Tensor>(0)->Shape(), context->Input<Tensor>(1)->Shape(), output_shape));
|
||||
for (int index = 2; index < input_count; index++) {
|
||||
TensorShape previous_output_shape = output_shape;
|
||||
TensorShape previous_output_shape = lhs_tensor->Shape();
|
||||
for (int index = 1; index < input_count; index++) {
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input<Tensor>(index)->Shape(), output_shape));
|
||||
previous_output_shape = output_shape;
|
||||
}
|
||||
Tensor* output_tensor = context->Output(0, output_shape);
|
||||
BinaryElementwisePreparation prepare(this);
|
||||
|
||||
auto rhs_tensor = context->Input<Tensor>(1);
|
||||
if (input_count == 2) {
|
||||
// special case for 2 tensors to avoid memset zero
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(context->Input<Tensor>(0), context->Input<Tensor>(1), output_tensor, &prepare));
|
||||
Impl_Add<CudaT>(
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(lhs_tensor, rhs_tensor, output_tensor, &prepare));
|
||||
ORT_RETURN_IF_ERROR(prepare.CopyToGpu());
|
||||
Impl_Compute(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
|
|
@ -263,9 +267,24 @@ Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
} else {
|
||||
// for more than 2 inputs, we need to accumulate into output tensor, as the shape from input0 + input1 might be different from output shape
|
||||
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
|
||||
for (int index = 0; index < input_count; index++) {
|
||||
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, lhs_tensor, output_tensor, &prepare));
|
||||
ORT_RETURN_IF_ERROR(prepare.CopyToGpu());
|
||||
Impl_Add(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
prepare.rhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.rhs_tensor->template Data<T>()),
|
||||
prepare.fdm_output_strides.GpuPtr(),
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size());
|
||||
|
||||
for (int index = 1; index < input_count; index++) {
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
|
||||
Impl_Add<CudaT>(
|
||||
Impl_Compute(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
|
|
@ -281,133 +300,34 @@ Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
|
|||
}
|
||||
return Status::OK();
|
||||
}
|
||||
template <typename T>
|
||||
Status Sum<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
|
||||
this->ComputeMethod(context, &Impl_Add);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status Max<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
typedef typename ToCudaType<T>::MappedType CudaT;
|
||||
const auto& node = Node();
|
||||
const auto& node_name = node.Name();
|
||||
auto input_count = node.InputArgCount().front();
|
||||
ORT_RETURN_IF_NOT(input_count >= 1, "Must have 1 or more inputs");
|
||||
|
||||
if (input_count == 1) {
|
||||
auto input_tensor = context->Input<Tensor>(0);
|
||||
const auto& input_shape = input_tensor->Shape();
|
||||
auto output_tensor = context->Output(0, input_shape);
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice));
|
||||
} else {
|
||||
// compute output shape first, using broadcast rule
|
||||
TensorShape output_shape;
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input<Tensor>(0)->Shape(), context->Input<Tensor>(1)->Shape(), output_shape));
|
||||
for (int index = 2; index < input_count; index++) {
|
||||
TensorShape previous_output_shape = output_shape;
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input<Tensor>(index)->Shape(), output_shape));
|
||||
}
|
||||
Tensor* output_tensor = context->Output(0, output_shape);
|
||||
BinaryElementwisePreparation prepare(this);
|
||||
this->ComputeMethod(context, &Impl_Max);
|
||||
|
||||
// More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly
|
||||
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
|
||||
Impl_Add<CudaT>(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
prepare.rhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.rhs_tensor->template Data<T>()),
|
||||
prepare.fdm_output_strides.GpuPtr(),
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size());
|
||||
for (int index = 1; index < input_count; index++) {
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
|
||||
Impl_Max<CudaT>(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
prepare.rhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.rhs_tensor->template Data<T>()),
|
||||
prepare.fdm_output_strides.GpuPtr(),
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size());
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
Status Min<T>::ComputeInternal(OpKernelContext* context) const {
|
||||
typedef typename ToCudaType<T>::MappedType CudaT;
|
||||
const auto& node = Node();
|
||||
const auto& node_name = node.Name();
|
||||
auto input_count = node.InputArgCount().front();
|
||||
ORT_RETURN_IF_NOT(input_count >= 1, "Must have 1 or more inputs");
|
||||
|
||||
if (input_count == 1) {
|
||||
auto input_tensor = context->Input<Tensor>(0);
|
||||
const auto& input_shape = input_tensor->Shape();
|
||||
auto output_tensor = context->Output(0, input_shape);
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice));
|
||||
} else {
|
||||
// compute output shape first, using broadcast rule
|
||||
TensorShape output_shape;
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input<Tensor>(0)->Shape(), context->Input<Tensor>(1)->Shape(), output_shape));
|
||||
for (int index = 2; index < input_count; index++) {
|
||||
TensorShape previous_output_shape = output_shape;
|
||||
ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input<Tensor>(index)->Shape(), output_shape));
|
||||
}
|
||||
Tensor* output_tensor = context->Output(0, output_shape);
|
||||
BinaryElementwisePreparation prepare(this);
|
||||
this->ComputeMethod(context, &Impl_Min);
|
||||
|
||||
// More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly
|
||||
CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT)));
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(0), output_tensor, &prepare));
|
||||
Impl_Add<CudaT>(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
prepare.rhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.rhs_tensor->template Data<T>()),
|
||||
prepare.fdm_output_strides.GpuPtr(),
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size());
|
||||
for (int index = 1; index < input_count; index++) {
|
||||
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input<Tensor>(index), output_tensor, &prepare));
|
||||
Impl_Min<CudaT>(
|
||||
prepare.output_rank_or_simple_broadcast,
|
||||
prepare.lhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.lhs_tensor->template Data<T>()),
|
||||
prepare.rhs_padded_strides.GpuPtr(),
|
||||
reinterpret_cast<const CudaT*>(prepare.rhs_tensor->template Data<T>()),
|
||||
prepare.fdm_output_strides.GpuPtr(),
|
||||
prepare.fdm_H,
|
||||
prepare.fdm_C,
|
||||
reinterpret_cast<CudaT*>(prepare.output_tensor->template MutableData<T>()),
|
||||
prepare.output_tensor->Shape().Size());
|
||||
}
|
||||
}
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
//Greater op output tensor type is bool, so it cannot directly fit in the macros
|
||||
//for other elementwise ops
|
||||
template <typename T, typename CudaT>
|
||||
Status CompareFunction<T, CudaT>::CompareMethod(OpKernelContext* context, void (*Impl_Compare)(
|
||||
size_t output_rank_or_simple_broadcast,
|
||||
const int64_t* lhs_padded_strides,
|
||||
const CudaT* lhs_data,
|
||||
const int64_t* rhs_padded_strides,
|
||||
const CudaT* rhs_data,
|
||||
const fast_divmod* fdm_output_strides,
|
||||
const fast_divmod& fdm_H,
|
||||
const fast_divmod& fdm_C,
|
||||
CudaT* output_data,
|
||||
size_t count)) const {
|
||||
Status CompareFunction<T, CudaT>::CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const {
|
||||
BinaryElementwisePreparation prepare(this);
|
||||
Prepare(context, &prepare);
|
||||
size_t output_size = prepare.output_tensor->Shape().Size();
|
||||
|
|
|
|||
|
|
@ -186,31 +186,50 @@ class PRelu final : public BinaryElementwise<ShouldBroadcast> {
|
|||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
||||
template <typename T, typename CudaT>
|
||||
class VariadicInputBase : public CudaKernel {
|
||||
public:
|
||||
VariadicInputBase(const OpKernelInfo& info) : CudaKernel(info) {}
|
||||
|
||||
Status ComputeInternal(OpKernelContext*) const override {
|
||||
return Status(common::ONNXRUNTIME, common::FAIL); // should not reach here
|
||||
}
|
||||
|
||||
typedef void (*ImplCompute)(size_t output_rank_or_simple_broadcast,
|
||||
const int64_t* lhs_padded_strides,
|
||||
const CudaT* lhs_data,
|
||||
const int64_t* rhs_padded_strides,
|
||||
const CudaT* rhs_data,
|
||||
const fast_divmod* fdm_output_strides,
|
||||
const fast_divmod& fdm_H,
|
||||
const fast_divmod& fdm_C,
|
||||
CudaT* output_data,
|
||||
size_t count);
|
||||
|
||||
Status ComputeMethod(OpKernelContext* context, ImplCompute Impl_Compute) const;
|
||||
};
|
||||
|
||||
// Sum allows varadic inputs, and it uses binary elementwise Add in implementation
|
||||
template <typename T>
|
||||
class Sum final : public CudaKernel {
|
||||
class Sum final : public VariadicInputBase<T, typename ToCudaType<T>::MappedType> {
|
||||
public:
|
||||
Sum(const OpKernelInfo& info) : CudaKernel(info) {
|
||||
}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
||||
|
||||
template <typename T>
|
||||
class Max final : public CudaKernel {
|
||||
public:
|
||||
Max(const OpKernelInfo& info) : CudaKernel(info) {
|
||||
}
|
||||
Sum(const OpKernelInfo& info) : VariadicInputBase<T, typename ToCudaType<T>::MappedType>(info) {}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class Min final : public CudaKernel {
|
||||
class Max final : public VariadicInputBase<T, typename ToCudaType<T>::MappedType> {
|
||||
public:
|
||||
Min(const OpKernelInfo& info) : CudaKernel(info) {
|
||||
}
|
||||
Max(const OpKernelInfo& info) : VariadicInputBase<T, typename ToCudaType<T>::MappedType>(info) {}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class Min final : public VariadicInputBase<T, typename ToCudaType<T>::MappedType> {
|
||||
public:
|
||||
Min(const OpKernelInfo& info) : VariadicInputBase<T, typename ToCudaType<T>::MappedType>(info) {}
|
||||
|
||||
Status ComputeInternal(OpKernelContext* context) const override;
|
||||
};
|
||||
|
|
@ -220,17 +239,18 @@ class CompareFunction : public BinaryElementwise<ShouldBroadcast> {
|
|||
public:
|
||||
CompareFunction(const OpKernelInfo& info) : BinaryElementwise(info) {}
|
||||
|
||||
Status CompareMethod(OpKernelContext* context, void (*Impl_Compare)(
|
||||
size_t output_rank_or_simple_broadcast,
|
||||
const int64_t* lhs_padded_strides,
|
||||
const CudaT* lhs_data,
|
||||
const int64_t* rhs_padded_strides,
|
||||
const CudaT* rhs_data,
|
||||
const fast_divmod* fdm_output_strides,
|
||||
const fast_divmod& fdm_H,
|
||||
const fast_divmod& fdm_C,
|
||||
CudaT* output_data,
|
||||
size_t count)) const;
|
||||
typedef void (*ImplCompare)(size_t output_rank_or_simple_broadcast,
|
||||
const int64_t* lhs_padded_strides,
|
||||
const CudaT* lhs_data,
|
||||
const int64_t* rhs_padded_strides,
|
||||
const CudaT* rhs_data,
|
||||
const fast_divmod* fdm_output_strides,
|
||||
const fast_divmod& fdm_H,
|
||||
const fast_divmod& fdm_C,
|
||||
CudaT* output_data,
|
||||
size_t count);
|
||||
|
||||
Status CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
|
|
|
|||
Loading…
Reference in a new issue