From e612018127358fa964da32d78cb082440771fe0d Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Thu, 7 Apr 2022 14:59:54 +0800 Subject: [PATCH] [CUDA] Tile Kernel Optimization (#11053) * tile cuda kernel optimization * resolve comments and fix win build error --- .../core/providers/cuda/tensor/tile.cc | 141 ++++------ .../core/providers/cuda/tensor/tile_impl.cu | 260 ++++++++++-------- .../core/providers/cuda/tensor/tile_impl.h | 30 +- .../test/providers/cpu/tensor/tile_op_test.cc | 5 + 4 files changed, 207 insertions(+), 229 deletions(-) diff --git a/onnxruntime/core/providers/cuda/tensor/tile.cc b/onnxruntime/core/providers/cuda/tensor/tile.cc index b7815a5337..350e628b78 100644 --- a/onnxruntime/core/providers/cuda/tensor/tile.cc +++ b/onnxruntime/core/providers/cuda/tensor/tile.cc @@ -40,6 +40,28 @@ ONNX_OPERATOR_KERNEL_EX( .TypeConstraint("T1", DataTypeImpl::GetTensorType()), Tile); +#define CASE_TILE(type) \ + case sizeof(type): { \ + TileImpl(Stream(), rank, fdm_input_shape, input_strides, \ + reinterpret_cast::MappedType*>(input_data), fdm_output_strides, \ + reinterpret_cast::MappedType*>(output_data), output_tensor.Shape().Size()); \ + } break + +#define CASE_TILE_MEMCPY(type) \ + case sizeof(type): { \ + TileMemcpyImpl(Stream(), reinterpret_cast::MappedType*>(input_data), \ + reinterpret_cast::MappedType*>(output_data), input_shape.Size(), \ + num_of_copies_per_batch); \ + } break + +#define CASE_TILE_BATCHED_MEMCPY(type) \ + case sizeof(type): { \ + TileBatchedMemcpyImpl(Stream(), reinterpret_cast::MappedType*>(input_data), \ + reinterpret_cast::MappedType*>(output_data), \ + num_of_elements_per_batch, input_shape.Size(), num_of_batch_copies, \ + num_of_copies_per_batch); \ + } break + Status Tile::ComputeInternal(OpKernelContext* ctx) const { auto& input_tensor = *ctx->Input(0); auto& repeats_tensor = *ctx->Input(1); @@ -62,6 +84,7 @@ Status Tile::ComputeInternal(OpKernelContext* ctx) const { void* output_data = output_tensor.MutableDataRaw(); const void* input_data = input_tensor.DataRaw(); + const auto element_size = input_tensor.DataType()->Size(); // Repeat tensor input can have 0 as a valid value // check if the computed output_shape size is 0 and @@ -88,66 +111,22 @@ Status Tile::ComputeInternal(OpKernelContext* ctx) const { num_of_copies_per_batch, num_of_batch_copies)) { if (!is_batched_memcpy) { - if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - input_shape.Size(), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - input_shape.Size(), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else if (input_tensor.IsDataType()) { - TileMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - input_shape.Size(), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else { - // Won't hit this as the kernel doesn't claim support for any type that will trigger this - ORT_THROW("Tile doesn't have an implementation yet for the type: ", input_tensor.DataType()); + switch (element_size) { + CASE_TILE_MEMCPY(float); + CASE_TILE_MEMCPY(double); + CASE_TILE_MEMCPY(MLFloat16); + default: + ORT_THROW("Unsupported value attribute datatype with sizeof=: ", element_size); + break; } } else { - if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileBatchedMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - num_of_elements_per_batch, - input_shape[0], // The tensor is atleast 1-D- this is safe - fast_divmod(static_cast(num_of_elements_per_batch * num_of_copies_per_batch)), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileBatchedMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - num_of_elements_per_batch, - input_shape[0], // The tensor is atleast 1-D- this is safe - fast_divmod(static_cast(num_of_elements_per_batch * num_of_copies_per_batch)), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else if (input_tensor.IsDataType()) { - TileBatchedMemcpyImpl( - Stream(), - reinterpret_cast::MappedType*>(input_data), - num_of_elements_per_batch, - input_shape[0], // The tensor is atleast 1-D- this is safe - fast_divmod(static_cast(num_of_elements_per_batch * num_of_copies_per_batch)), - reinterpret_cast::MappedType*>(output_data), - output_shape.Size()); - } else { - // Won't hit this as the kernel doesn't claim support for any type that will trigger this - ORT_THROW("Tile doesn't have an implementation yet for the type: ", input_tensor.DataType()); + switch (element_size) { + CASE_TILE_BATCHED_MEMCPY(float); + CASE_TILE_BATCHED_MEMCPY(double); + CASE_TILE_BATCHED_MEMCPY(MLFloat16); + default: + ORT_THROW("Unsupported value attribute datatype with sizeof=: ", element_size); + break; } } @@ -158,8 +137,8 @@ Status Tile::ComputeInternal(OpKernelContext* ctx) const { TArray input_strides(input_pitches); TArray fdm_input_shape(rank); - for (int32_t i = 0; i < input_dims.size(); ++i) { - fdm_input_shape[i] = fast_divmod(gsl::narrow_cast(input_dims[i])); + for (size_t i = 0; i < input_dims.size(); ++i) { + fdm_input_shape[gsl::narrow_cast(i)] = fast_divmod(gsl::narrow_cast(input_dims[i])); } TArray fdm_output_strides(rank); @@ -172,41 +151,13 @@ Status Tile::ComputeInternal(OpKernelContext* ctx) const { static_assert(sizeof(double) == sizeof(int64_t), "Double and Int64 are of different sizes"); if (output_tensor.Shape().Size() > 0) { - if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileImpl( - Stream(), - rank, - fdm_input_shape, - input_strides, - reinterpret_cast::MappedType*>(input_data), - fdm_output_strides, - reinterpret_cast::MappedType*>(output_data), - output_tensor.Shape().Size()); - } else if (input_tensor.IsDataType() || - input_tensor.IsDataType()) { - TileImpl( - Stream(), - rank, - fdm_input_shape, - input_strides, - reinterpret_cast::MappedType*>(input_data), - fdm_output_strides, - reinterpret_cast::MappedType*>(output_data), - output_tensor.Shape().Size()); - } else if (input_tensor.IsDataType()) { - TileImpl( - Stream(), - rank, - fdm_input_shape, - input_strides, - reinterpret_cast::MappedType*>(input_data), - fdm_output_strides, - reinterpret_cast::MappedType*>(output_data), - output_tensor.Shape().Size()); - } else { - // Won't hit this as the kernel doesn't claim support for any type that will trigger this - ORT_THROW("Tile doesn't have an implementation yet for the type: ", input_tensor.DataType()); + switch (element_size) { + CASE_TILE(float); + CASE_TILE(double); + CASE_TILE(MLFloat16); + default: + ORT_THROW("Unsupported value attribute datatype with sizeof=: ", element_size); + break; } } diff --git a/onnxruntime/core/providers/cuda/tensor/tile_impl.cu b/onnxruntime/core/providers/cuda/tensor/tile_impl.cu index adbf095c55..29632444b7 100644 --- a/onnxruntime/core/providers/cuda/tensor/tile_impl.cu +++ b/onnxruntime/core/providers/cuda/tensor/tile_impl.cu @@ -7,140 +7,178 @@ namespace onnxruntime { namespace cuda { -constexpr int MAX_DIMS = 16; +#ifdef USE_ROCM +constexpr int num_elements_per_thread = 2; +constexpr int num_threads_per_block = 512; +#else +constexpr int num_elements_per_thread = GridDim::maxElementsPerThread; +constexpr int num_threads_per_block = GridDim::maxThreadsPerBlock; +#endif template -__global__ void _UnRolledTileKernel( - const size_t shape_rank, - const TArray fdm_input_shape, - const TArray input_strides, - const T* input_data, - const TArray fdm_output_strides, - T* output_data, - const CUDA_LONG N) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - CUDA_LONG input_index = 0; - CUDA_LONG output_index = id; +__global__ void _UnRolledTileKernel(const size_t shape_rank, const TArray fdm_input_shape, + const TArray input_strides, const T* input_data, + const TArray fdm_output_strides, T* output_data, const CUDA_LONG N) { + CUDA_LONG start = num_elements_per_thread * num_threads_per_block * blockIdx.x + threadIdx.x; + T value[num_elements_per_thread]; + CUDA_LONG id = start; +#pragma unroll + for (int i = 0; i < num_elements_per_thread; ++i) { + if (id < N) { + CUDA_LONG input_index = 0; + CUDA_LONG offset = id; +#pragma unroll + for (auto dim = 0; dim < fdm_output_strides.Capacity(); ++dim) { + if (dim >= shape_rank) { + break; + } - #pragma unroll - for (int dim = 0; dim < MAX_DIMS; ++dim) { - if (dim == shape_rank) { - break; + int q, r; + fdm_output_strides[dim].divmod(offset, q, r); + int in_coord = fdm_input_shape[dim].mod(q); + input_index += input_strides[dim] * in_coord; + offset = r; + } + + value[i] = input_data[input_index]; + id += num_threads_per_block; } - int out_coord, r; - fdm_output_strides[dim].divmod(output_index, out_coord, r); - output_index = r; - int in_coord = fdm_input_shape[dim].mod(out_coord); - input_index += input_strides[dim] * in_coord; } - output_data[id] = input_data[input_index]; -} -template -__global__ void _TileKernel( - const size_t shape_rank, - const TArray fdm_input_shape, - const TArray input_strides, - const T* input_data, - const TArray fdm_output_strides, - T* output_data, - const CUDA_LONG N) { - CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - CUDA_LONG input_index = 0; - CUDA_LONG output_index = id; - for (int dim = 0; dim < shape_rank; ++dim) { - int out_coord, r; - fdm_output_strides[dim].divmod(output_index, out_coord, r); - output_index = r; - int in_coord = fdm_input_shape[dim].mod(out_coord); - input_index += input_strides[dim] * in_coord; - } - output_data[id] = input_data[input_index]; -} - -template -void TileImpl( - cudaStream_t stream, - const size_t shape_rank, - const TArray& fdm_input_shape, - const TArray& input_stride, - const T* input_data, - const TArray& fdm_output_strides, - T* output_data, - const size_t N) { - int blocksPerGrid = (int)(ceil(static_cast(N) / GridDim::maxThreadsPerBlock)); - if (shape_rank > MAX_DIMS) { - _TileKernel<<>>( - shape_rank, fdm_input_shape, input_stride, input_data, - fdm_output_strides, output_data, (CUDA_LONG)N); - } else { - _UnRolledTileKernel<<>>( - shape_rank, fdm_input_shape, input_stride, input_data, - fdm_output_strides, output_data, (CUDA_LONG)N); + id = start; +#pragma unroll + for (int i = 0; i < num_elements_per_thread; ++i) { + if (id < N) { + output_data[id] = value[i]; + id += num_threads_per_block; + } } } template -__global__ void _TileMemcpyKernel( - const T* input_data, - const size_t num_input_elements, - T* output_data, - const size_t N) { +void TileImpl(cudaStream_t stream, const size_t shape_rank, const TArray& fdm_input_shape, + const TArray& input_stride, const T* input_data, const TArray& fdm_output_strides, + T* output_data, const size_t N) { + int blocksPerGrid = static_cast(CeilDiv(N, num_threads_per_block * num_elements_per_thread)); + _UnRolledTileKernel<<>>(shape_rank, fdm_input_shape, input_stride, + input_data, fdm_output_strides, + output_data, static_cast(N)); +} + +template +__global__ void _TileMemcpyKernel(const T* input_data, T* output_data, const CUDA_LONG N, const size_t repeats) { CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - auto input_index = id % num_input_elements; - output_data[id] = input_data[input_index]; + T input_val = input_data[id]; + for (size_t i = 0; i < repeats; ++i) { + output_data[id] = input_val; + id += N; + } } template -void TileMemcpyImpl( - cudaStream_t stream, - const T* input_data, - const size_t num_input_elements, - T* output_data, - const size_t num_output_elements) { - int blocksPerGrid = (int)(ceil(static_cast(num_output_elements) / GridDim::maxThreadsPerBlock)); - _TileMemcpyKernel<<>>( - input_data, num_input_elements, output_data, (CUDA_LONG)num_output_elements); +void TileMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, const size_t num_input_elements, + const size_t repeats) { + using Vec4T = aligned_vector; + using Vec2T = aligned_vector; + constexpr int vec4_alignment = std::alignment_of::value; + constexpr int vec2_alignment = std::alignment_of::value; + uint64_t address_input = reinterpret_cast(input_data); + uint64_t address_output = reinterpret_cast(output_data); + CUDA_LONG N = static_cast(num_input_elements); + if (num_input_elements % 4 == 0 && address_input % vec4_alignment == 0 && address_output % vec4_alignment == 0) { + N /= 4; + int blocksPerGrid = CeilDiv(N, num_threads_per_block); + _TileMemcpyKernel<<>>( + reinterpret_cast(input_data), reinterpret_cast(output_data), N, repeats); + return; + } else if (num_input_elements % 2 == 0 && address_input % vec2_alignment == 0 && + address_output % vec2_alignment == 0) { + N /= 2; + int blocksPerGrid = CeilDiv(N, num_threads_per_block); + _TileMemcpyKernel<<>>( + reinterpret_cast(input_data), reinterpret_cast(output_data), N, repeats); + return; + } + + int blocksPerGrid = CeilDiv(N, num_threads_per_block); + _TileMemcpyKernel<<>>(input_data, output_data, N, repeats); } +// Input size is [batch, data], output size is [batch * batch_repeats, data * repeats_per_batch]. +// Here size_input_row = data, size_output_row = data * repeats_per_batch, +// size_output_batch = batch * data * repeats_per_batch template -__global__ void _TileBatchedMemcpyKernel( - const T* input_data, - const size_t num_of_elements_per_input_batch, - const size_t num_input_batch_count, - const fast_divmod num_of_elements_per_output_batch, - T* output_data, - const size_t N) { +__global__ void _TileBatchedMemcpyKernel(const T* input_data, T* output_data, const fast_divmod divmod_size_input_row, + const CUDA_LONG size_input_row, const CUDA_LONG size_output_row, + const CUDA_LONG size_output_batch, const size_t batch_repeats, + const size_t repeats_per_batch, const CUDA_LONG N) { CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(id, N); - CUDA_LONG batch_idx = 0; - CUDA_LONG element_idx = 0; - num_of_elements_per_output_batch.divmod(id, batch_idx, element_idx); - output_data[id] = input_data[(batch_idx % num_input_batch_count) * num_of_elements_per_input_batch + (element_idx % num_of_elements_per_input_batch)]; + T input_val = input_data[id]; + CUDA_LONG q, r; + divmod_size_input_row.divmod(id, q, r); + CUDA_LONG batch_offset = q * size_output_row + r; + for (size_t i = 0; i < batch_repeats; ++i) { + CUDA_LONG offset = batch_offset; + for (size_t j = 0; j < repeats_per_batch; ++j) { + output_data[offset] = input_val; + offset += size_input_row; + } + batch_offset += size_output_batch; + } } +// Input size is [batch, data], output size is [batch * batch_repeats, data * repeats_per_batch]. +// Here size_input_row = data, num_input_elements = batch * data template -void TileBatchedMemcpyImpl( - cudaStream_t stream, - const T* input_data, - const size_t num_of_elements_per_input_batch, - const size_t num_input_batch_count, - const fast_divmod& num_of_elements_per_output_batch, - T* output_data, - const size_t num_output_elements) { - int blocksPerGrid = (int)(ceil(static_cast(num_output_elements) / GridDim::maxThreadsPerBlock)); - _TileBatchedMemcpyKernel<<>>( - input_data, - num_of_elements_per_input_batch, - num_input_batch_count, - num_of_elements_per_output_batch, - output_data, - (CUDA_LONG)num_output_elements); +void TileBatchedMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, const size_t size_input_row, + const size_t num_input_elements, const size_t batch_repeats, + const size_t repeats_per_batch) { + using Vec4T = aligned_vector; + using Vec2T = aligned_vector; + constexpr int vec4_alignment = std::alignment_of::value; + constexpr int vec2_alignment = std::alignment_of::value; + uint64_t address_input = reinterpret_cast(input_data); + uint64_t address_output = reinterpret_cast(output_data); + CUDA_LONG size_input_row_vec = static_cast(size_input_row); + CUDA_LONG N = static_cast(num_input_elements); + if (size_input_row % 4 == 0 && address_input % vec4_alignment == 0 && address_output % vec4_alignment == 0) { + size_input_row_vec /= 4; + N /= 4; + int blocksPerGrid = CeilDiv(N, num_threads_per_block); + _TileBatchedMemcpyKernel<<>>( + reinterpret_cast(input_data), reinterpret_cast(output_data), + fast_divmod(size_input_row_vec), size_input_row_vec, + size_input_row_vec * static_cast(repeats_per_batch), N * static_cast(repeats_per_batch), + batch_repeats, repeats_per_batch, N); + return; + } else if (size_input_row % 2 == 0 && address_input % vec2_alignment == 0 && address_output % vec2_alignment == 0) { + size_input_row_vec /= 2; + N /= 2; + int blocksPerGrid = CeilDiv(N, num_threads_per_block); + _TileBatchedMemcpyKernel<<>>( + reinterpret_cast(input_data), reinterpret_cast(output_data), + fast_divmod(size_input_row_vec), size_input_row_vec, + size_input_row_vec * static_cast(repeats_per_batch), N * static_cast(repeats_per_batch), + batch_repeats, repeats_per_batch, N); + return; + } + + int blocksPerGrid = static_cast(CeilDiv(N, num_threads_per_block)); + _TileBatchedMemcpyKernel<<>>( + input_data, output_data, fast_divmod(size_input_row_vec), size_input_row_vec, + size_input_row_vec * static_cast(repeats_per_batch), N * static_cast(repeats_per_batch), + batch_repeats, repeats_per_batch, N); } -#define SPECIALIZED_IMPL(T) \ - template void TileImpl(cudaStream_t stream, const size_t shape_rank, const TArray& fdm_input_shape, const TArray& input_stride, const T* input_data, const TArray& fdm_output_strides, T* output_data, const size_t N); \ - template void TileMemcpyImpl(cudaStream_t stream, const T* input_data, const size_t num_input_elements, T* output_data, const size_t num_output_elements); \ - template void TileBatchedMemcpyImpl(cudaStream_t stream, const T* input_data, const size_t num_of_elements_per_input_batch, const size_t num_input_batch_count, const fast_divmod& num_of_elements_per_output_batch, T* output_data, const size_t num_output_elements); +#define SPECIALIZED_IMPL(T) \ + template void TileImpl(cudaStream_t stream, const size_t shape_rank, const TArray& fdm_input_shape, \ + const TArray& input_stride, const T* input_data, \ + const TArray& fdm_output_strides, T* output_data, const size_t N); \ + template void TileMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, \ + const size_t num_input_elements, const size_t repeats); \ + template void TileBatchedMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, \ + const size_t size_input_row, const size_t num_input_elements, \ + const size_t batch_repeats, const size_t repeats_per_batch); SPECIALIZED_IMPL(float) SPECIALIZED_IMPL(double) diff --git a/onnxruntime/core/providers/cuda/tensor/tile_impl.h b/onnxruntime/core/providers/cuda/tensor/tile_impl.h index a612beabcf..6f2d56c6c8 100644 --- a/onnxruntime/core/providers/cuda/tensor/tile_impl.h +++ b/onnxruntime/core/providers/cuda/tensor/tile_impl.h @@ -9,33 +9,17 @@ namespace onnxruntime { namespace cuda { template -void TileImpl( - cudaStream_t stream, - const size_t shape_rank, - const TArray& input_shape, - const TArray& input_strides, - const T* input_data, - const TArray& fdm_output_strides, - T* output_data, - const size_t N); +void TileImpl(cudaStream_t stream, const size_t shape_rank, const TArray& fdm_input_shape, + const TArray& input_stride, const T* input_data, const TArray& fdm_output_strides, + T* output_data, const size_t N); template -void TileMemcpyImpl( - cudaStream_t stream, - const T* input_data, - const size_t num_input_elements, - T* output_data, - const size_t num_output_elements); +void TileMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, const size_t num_input_elements, + const size_t repeats); template -void TileBatchedMemcpyImpl( - cudaStream_t stream, - const T* input_data, - const size_t num_of_elements_per_input_batch, - const size_t num_input_batch_count, - const fast_divmod& num_of_elements_per_output_batch, - T* output_data, - const size_t num_output_elements); +void TileBatchedMemcpyImpl(cudaStream_t stream, const T* input_data, T* output_data, const size_t size_input_row, + const size_t num_input_elements, const size_t batch_repeats, const size_t repeats_per_batch); } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/tile_op_test.cc b/onnxruntime/test/providers/cpu/tensor/tile_op_test.cc index f992992c9d..45c36417f5 100644 --- a/onnxruntime/test/providers/cpu/tensor/tile_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/tile_op_test.cc @@ -37,9 +37,14 @@ void RunTestWrapper() { // Tile2D_1Axis RunTest({11, 12, 21, 22}, {2, 2}, {2, 1}, {2}, {11, 12, 21, 22, 11, 12, 21, 22}, {4, 2}); + RunTest({11, 12, 21, 22, 31, 32}, {2, 3}, {2, 1}, {2}, {11, 12, 21, 22, 31, 32, 11, 12, 21, 22, 31, 32}, {4, 3}); // Tile2D_2Axes RunTest({11, 12, 21, 22}, {2, 2}, {2, 2}, {2}, {11, 12, 11, 12, 21, 22, 21, 22, 11, 12, 11, 12, 21, 22, 21, 22}, {4, 4}); + RunTest({11, 12, 13, 14, 21, 22, 23, 24}, {2, 4}, {2, 2}, {2}, + {11, 12, 13, 14, 11, 12, 13, 14, 21, 22, 23, 24, 21, 22, 23, 24, + 11, 12, 13, 14, 11, 12, 13, 14, 21, 22, 23, 24, 21, 22, 23, 24}, + {4, 8}); // Tile3D RunTest({111, 112, 113, 122, 123, 124}, {2, 1, 3}, {1, 2, 1}, {3}, {111, 112, 113, 111, 112, 113, 122, 123, 124, 122, 123, 124}, {2, 2, 3});