saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make ROCM and CUDA reduction_all code more similar.

Browse source
This commit is contained in:
Jesse Benson 2020-12-09 12:59:24 -08:00 committed by Jesse Benson
parent 1eb146f561
commit 38c49c2483
4 changed files with 24 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
onnxruntime/core/providers/rocm/reduction/reduction_utils.cuh
View file

@ -54,15 +54,5 @@ struct Identity {
}
};
template <typename T>
struct ToBuffer {
typedef T Type;
};
template <>
struct ToBuffer<half> {
typedef float Type;
};
} // namespace rocm
} // namespace onnxruntime
} // namespace onnxruntime

7
orttraining/orttraining/training_ops/cuda/reduction/reduction_all.cu
View file

@ -14,18 +14,19 @@ namespace cuda {
template <typename Tin, typename Tout>
__global__ void ScalarSqrtKernel(Tin* input, Tout* output) {
*output = (Tout)_Sqrt(*input);
};
}
template <typename Tin, typename Tout>
void ScalarSqrt(Tin* input, Tout* output) {
ScalarSqrtKernel<<<1, 1, 0>>>(input, output);
};
}
template void ScalarSqrt(float* input, float* output);
template void ScalarSqrt(half* input, half* output);
template void ScalarSqrt(float* input, half* output);
template <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
__launch_bounds__(ChunkGroup<1>::thread_count_per_block)
__global__ void MultiTensorReduceKernel(ChunkGroup<1> chunk_group, TOut* output) {
const int group_index = chunk_group.block_index_to_tensor_group_index[blockIdx.x];
const int tensor_size = chunk_group.tensor_sizes[group_index];
@ -79,7 +80,7 @@ __global__ void MultiTensorReduceKernel(ChunkGroup<1> chunk_group, TOut* output)
if (threadIdx.x == 0) {
atomic_add(w_norm, TOutOp()(TOut(shared_memory[0])));
}
};
}
template <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) {

28
orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cc
View file

@ -2,23 +2,14 @@
// Licensed under the MIT License.
#include "orttraining/training_ops/rocm/reduction/reduction_all.h"
#include "core/providers/rocm/reduction/reduction_functions.h"
#include "core/framework/op_kernel_context_internal.h"
#include "core/providers/rocm/reduction/reduction_functions.h"
#include "core/providers/rocm/shared_inc/accumulation_type.h"
namespace onnxruntime {
namespace rocm {
template <typename T>
struct AccumulateType {};
template <>
struct AccumulateType<float> { using type = float; };
template <>
struct AccumulateType<half> { using type = float; };
template <>
struct AccumulateType<double> { using type = double; };
template <typename T>
using AccType = typename AccumulateType<T>::type;
#define REGISTER_REDUCE_ALL_KERNEL_TYPED(Name, TIn, TOut) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
Name, \
@ -58,7 +49,6 @@ Status ReduceAllL2<TIn, TOut>::ComputeInternal(OpKernelContext* ctx) const {
// bool deterministic = ctx->GetUseDeterministicCompute();
bool deterministic = true;
if (!deterministic) {
typedef MultiTensorReduceL2<HipTIn, HipTOut> TFunctor;
TFunctor functor;
@ -70,11 +60,9 @@ Status ReduceAllL2<TIn, TOut>::ComputeInternal(OpKernelContext* ctx) const {
// *p_output is the squared sum of all elements.
// Let's take a sqrt to get the actual L2-norm.
ScalarSqrt(p_output, p_output);
}
else {
} else {
// alternate path only for deterministic compute ..
typedef AccType<HipTOut> HipTAcc;
typedef AccumulationType_t<HipTOut> HipTAcc;
// find scratch buffer size needed by 'reduce_square_sum' for each tensor
int scratch_size = 0;
@ -86,16 +74,16 @@ Status ReduceAllL2<TIn, TOut>::ComputeInternal(OpKernelContext* ctx) const {
scratch_size = std::max(scratch_size, compute_reduction_buffer_size(sizeof(HipTAcc), total_tensor_count));
// add head room for final output and square norms of each tensor
scratch_size += (1 + total_tensor_count)*sizeof(HipTAcc);
scratch_size += (1 + total_tensor_count) * sizeof(HipTAcc);
// create GPU scratch space and zero target for each tensor square norm
auto scratch_buffer = GetScratchBuffer<uint8_t>(scratch_size);
HIP_RETURN_IF_ERROR(hipMemsetAsync(scratch_buffer.get(), 0, sizeof(HipTAcc)*(1 + total_tensor_count)));
HIP_RETURN_IF_ERROR(hipMemsetAsync(scratch_buffer.get(), 0, sizeof(HipTAcc) * (1 + total_tensor_count)));
HipTAcc* p_global_sqnorm = reinterpret_cast<HipTAcc*>(scratch_buffer.get());
HipTAcc* p_tensor_sqnorm = p_global_sqnorm + 1;
HipTAcc* p_reduce_buffer = p_tensor_sqnorm + total_tensor_count;
// perform reduction l2norm = sqrt[sum(tensor[i][j]**2)] for i,j over all tensor elements
for (int i = 0; i < total_tensor_count; ++i) {
HipTIn* p_tensor_i = reinterpret_cast<HipTIn*>(grouped_tensor_pointers[i][0]);

20
orttraining/orttraining/training_ops/rocm/reduction/reduction_all.cu
View file

@ -1,32 +1,34 @@
#include "hip/hip_runtime.h"
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <hip/hip_runtime.h>
#include "orttraining/training_ops/rocm/reduction/reduction_all.h"
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/rocm_common.h"
#include "core/providers/rocm/atomic/common.cuh"
#include "core/providers/rocm/reduction/reduction_utils.cuh"
#include "core/providers/rocm/shared_inc/accumulation_type.h"
namespace onnxruntime {
namespace rocm {
template<typename Tin, typename Tout>
__global__ void _ScalarSqrtImpl(Tin* input, Tout* output) {
__global__ void ScalarSqrtKernel(Tin* input, Tout* output) {
*output = (Tout)_Sqrt(*input);
};
}
template<typename Tin, typename Tout>
void ScalarSqrt(Tin* input, Tout* output) {
hipLaunchKernelGGL(_ScalarSqrtImpl, dim3(1), dim3(1), 0, 0, input, output);
};
hipLaunchKernelGGL(ScalarSqrtKernel, dim3(1), dim3(1), 0, 0, input, output);
}
template void ScalarSqrt(float* input, float* output);
template void ScalarSqrt(half* input, half* output);
template void ScalarSqrt(float* input, half* output);
template <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
__global__ void _MultiTensorReduceImpl(ChunkGroup<1> chunk_group, TOut* output) {
__launch_bounds__(ChunkGroup<1>::thread_count_per_block)
__global__ void MultiTensorReduceKernel(ChunkGroup<1> chunk_group, TOut* output) {
const int group_index = chunk_group.block_index_to_tensor_group_index[blockIdx.x];
const int tensor_size = chunk_group.tensor_sizes[group_index];
const int chunk_size = chunk_group.chunk_size;
@ -79,7 +81,7 @@ __global__ void _MultiTensorReduceImpl(ChunkGroup<1> chunk_group, TOut* output)
if (threadIdx.x == 0) {
atomic_add(w_norm, TOutOp()(shared_memory[0]));
}
};
}
template <typename TIn, typename TOut, typename TBuf, typename TInOp, typename TOutOp>
void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) {
@ -92,12 +94,12 @@ void MultiTensorReduce(ChunkGroup<1> chunk_group, TOut* output) {
ORT_ENFORCE(thread_count % GPU_WARP_SIZE == 0);
ORT_ENFORCE((thread_count & (thread_count - 1)) == 0);
hipLaunchKernelGGL(HIP_KERNEL_NAME(_MultiTensorReduceImpl<TIn, TOut, TBuf, TInOp, TOutOp>), dim3(chunk_group.chunk_count), dim3(thread_count), shared_memory_size, 0, chunk_group, output);
hipLaunchKernelGGL(HIP_KERNEL_NAME(MultiTensorReduceKernel<TIn, TOut, TBuf, TInOp, TOutOp>), dim3(chunk_group.chunk_count), dim3(thread_count), shared_memory_size, 0, chunk_group, output);
}
template <typename TIn, typename TOut>
void MultiTensorReduceL2<TIn, TOut>::operator()(ChunkGroup<1> chunk_group, TOut* output) {
typedef typename ToBuffer<TIn>::Type TBuf;
using TBuf = AccumulationType_t<TIn>;
MultiTensorReduce<TIn, TOut, TBuf, Square<TBuf, TIn>, Cast<TOut, TBuf>>(chunk_group, output);
}