refactor Cuda Ops Sum, Max, Min, remove dup code (#1946)

refactor Cuda Ops Sum, Max, Min, remove dup code
This commit is contained in:
Hector Li 2019-10-07 13:17:49 -07:00 committed by GitHub
parent 7b39f5090c
commit 00e24ae4fe
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 89 additions and 149 deletions

View file

@ -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();

View file

@ -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>