mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Adding Transpose3d and Transpose4d special case kernels for Rocm (#5837)
* add transpose3d; seeing memory fault on rocm3.7 * cleaned up code; commit to switch machines * tested working on gcr-openpai-35; 168 ex/sec * remove debug HCC_ENABLE_PRINTF
This commit is contained in:
Suffian Khan
GitHub
parent
b495ae8103
commit
7196d4206f
3 changed files with 282 additions and 7 deletions
|
|
@ -103,6 +103,53 @@ Status Transpose::DoTranspose(const Transpose& kernel,
|
|||
TensorPitches original_input_strides(input_dims);
|
||||
TensorPitches original_output_strides(output_dims);
|
||||
|
||||
// flatten the adjacent dimensions which are contiguous
|
||||
// for example: permutations[0, 2, 3, 1] -> [0, 2, 1], permutations[0, 3, 1, 2] -> [0, 2, 1]
|
||||
auto new_rank = rank;
|
||||
std::vector<size_t> new_permutations(permutations);
|
||||
std::vector<int64_t> new_input_dims(input_dims);
|
||||
std::vector<int64_t> new_output_dims(output_dims);
|
||||
|
||||
for (auto i = rank - 1; i > 0; i--) {
|
||||
auto curr = new_permutations[i];
|
||||
auto prev = new_permutations[i - 1];
|
||||
if (prev + 1 == curr) {
|
||||
// all dims bigger than curr need to be reduced by 1 due to the merging.
|
||||
for (auto j = 0; j < new_rank; j++) {
|
||||
if (new_permutations[j] > curr) {
|
||||
new_permutations[j] -= 1;
|
||||
}
|
||||
}
|
||||
for (auto j = i+1; j < new_rank; j++) {
|
||||
new_permutations[j-1] = new_permutations[j];
|
||||
}
|
||||
|
||||
// update input dims
|
||||
new_input_dims[prev] *= new_input_dims[curr];
|
||||
new_input_dims[curr] = 1;
|
||||
for (auto j = static_cast<int32_t>(curr+1); j < new_rank; j++) {
|
||||
new_input_dims[j-1] = new_input_dims[j];
|
||||
}
|
||||
new_input_dims[new_rank-1] = 1;
|
||||
|
||||
// update output dims
|
||||
new_output_dims[i-1] *= new_output_dims[i];
|
||||
new_output_dims[i] = 1;
|
||||
for (auto j = i+1; j < new_rank; j++) {
|
||||
new_output_dims[j-1] = new_output_dims[j];
|
||||
}
|
||||
new_output_dims[new_rank-1] = 1;
|
||||
|
||||
new_rank--;
|
||||
}
|
||||
}
|
||||
new_permutations.resize(new_rank);
|
||||
new_input_dims.resize(new_rank);
|
||||
new_output_dims.resize(new_rank);
|
||||
|
||||
TensorPitches new_input_strides(new_input_dims);
|
||||
TensorPitches new_output_strides(new_output_dims);
|
||||
|
||||
// TArray<int64_t> input_strides(rank);
|
||||
// for (auto i = 0; i < rank; i++) {
|
||||
// input_strides[i] = original_input_strides[permutations[i]];
|
||||
|
|
@ -113,13 +160,32 @@ Status Transpose::DoTranspose(const Transpose& kernel,
|
|||
// output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
|
||||
// }
|
||||
|
||||
RocmAsyncBuffer<int64_t> input_strides(&kernel, rank);
|
||||
for (auto i = 0; i < rank; i++) {
|
||||
input_strides.CpuPtr()[i] = original_input_strides[permutations[i]];
|
||||
size_t element_size = input.DataType()->Size();
|
||||
std::vector<int64_t> input_shape(new_rank);
|
||||
std::vector<int64_t> tmp_input_strides(new_rank);
|
||||
std::vector<int64_t> tmp_output_strides(new_rank);
|
||||
for (auto i = 0; i < new_rank; i++) {
|
||||
input_shape[i] = new_input_dims[i];
|
||||
tmp_input_strides[i] = new_input_strides[i];
|
||||
tmp_output_strides[i] = new_output_strides[new_permutations[i]];
|
||||
}
|
||||
|
||||
RocmAsyncBuffer<fast_divmod> output_strides(&kernel, rank);
|
||||
ORT_ENFORCE(CalculateFdmStrides(output_strides.CpuSpan(), output_dims));
|
||||
if (CanDoTranspose3D(new_rank, new_input_dims, new_permutations)) {
|
||||
return Transpose3DImpl(kernel, element_size, input_shape, tmp_input_strides,
|
||||
input.DataRaw(), output.MutableDataRaw(), output.Shape().Size());
|
||||
|
||||
} else if (CanDoTranspose4D(kernel.GetDeviceProp(), element_size, new_rank, new_input_dims, new_permutations)) {
|
||||
return Transpose4DImpl(kernel, element_size, input_shape, tmp_input_strides, input.DataRaw(),
|
||||
tmp_output_strides, output.MutableDataRaw(), output.Shape().Size());
|
||||
}
|
||||
|
||||
RocmAsyncBuffer<int64_t> input_strides(&kernel, new_rank);
|
||||
for (auto i = 0; i < new_rank; i++) {
|
||||
input_strides.CpuPtr()[i] = new_input_strides[new_permutations[i]];
|
||||
}
|
||||
|
||||
RocmAsyncBuffer<fast_divmod> output_strides(&kernel, new_rank);
|
||||
ORT_ENFORCE(CalculateFdmStrides(output_strides.CpuSpan(), new_output_dims));
|
||||
|
||||
// TODO: use output shape in reverse order for uint24 math
|
||||
// for (auto i = 0; i < rank; i++) {
|
||||
|
|
@ -131,8 +197,7 @@ Status Transpose::DoTranspose(const Transpose& kernel,
|
|||
ORT_RETURN_IF_ERROR(input_strides.CopyToGpu());
|
||||
ORT_RETURN_IF_ERROR(output_strides.CopyToGpu());
|
||||
|
||||
size_t element_size = input.DataType()->Size();
|
||||
auto status = TransposeImpl(element_size, rank, input_strides.GpuPtr(), input.DataRaw(),
|
||||
auto status = TransposeImpl(element_size, new_rank, input_strides.GpuPtr(), input.DataRaw(),
|
||||
output_strides.GpuPtr(), output.MutableDataRaw(), output.Shape().Size());
|
||||
|
||||
return status;
|
||||
|
|
|
|||
|
|
@ -7,6 +7,199 @@
|
|||
namespace onnxruntime {
|
||||
namespace rocm {
|
||||
|
||||
// todo: should it be different for AMD HW?
|
||||
constexpr int TILE_DIM = 16;
|
||||
|
||||
template <typename T>
|
||||
/* __global__ void Transpose3DKernel(const TArray<int64_t> input_shape,
|
||||
const TArray<int64_t> input_strides,
|
||||
const T* input_data, T* output_data) { */
|
||||
__global__ void Transpose3DKernel(const int64_t* input_shape,
|
||||
const int64_t* input_strides,
|
||||
const T* __restrict__ input_data, T* __restrict__ output_data) {
|
||||
|
||||
__shared__ T tile[TILE_DIM * (TILE_DIM + 1)];
|
||||
|
||||
int x = blockIdx.x * TILE_DIM + threadIdx.x;
|
||||
int y = blockIdx.y * TILE_DIM + threadIdx.y;
|
||||
|
||||
tile[threadIdx.y * TILE_DIM + threadIdx.x] = input_data[blockIdx.z * input_strides[0] + y * input_shape[2] + x];
|
||||
__syncthreads();
|
||||
|
||||
x = blockIdx.y * TILE_DIM + threadIdx.x;
|
||||
y = blockIdx.x * TILE_DIM + threadIdx.y;
|
||||
|
||||
output_data[blockIdx.z * input_strides[0] + y * input_shape[1] + x] = tile[threadIdx.x * TILE_DIM + threadIdx.y];
|
||||
}
|
||||
|
||||
bool CanDoTranspose3D(int32_t rank,
|
||||
const std::vector<int64_t>& input_dims,
|
||||
const std::vector<size_t>& permutations) {
|
||||
if (rank == 3 &&
|
||||
// permutation is done in the last two dimensions.
|
||||
permutations[rank - 2] == (rank - 1) && permutations[rank - 1] == (rank - 2) &&
|
||||
// the last two dimensions are aligned with TILE_DIM.
|
||||
input_dims[rank - 2] % TILE_DIM == 0 && input_dims[rank - 1] % TILE_DIM == 0) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/* Status Transpose3DImpl(size_t element_size,
|
||||
const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides,
|
||||
const void* input_data, void* output_data, int64_t N) { */
|
||||
Status Transpose3DImpl(const Transpose& kernel, size_t element_size,
|
||||
const std::vector<int64_t>& input_shape, const std::vector<int64_t>& input_strides,
|
||||
const void* input_data, void* output_data, int64_t N) {
|
||||
dim3 block_size(TILE_DIM, TILE_DIM);
|
||||
dim3 grid_size(input_shape[2] / TILE_DIM, input_shape[1] / TILE_DIM, input_shape[0]);
|
||||
|
||||
RocmKernel::RocmAsyncBuffer<int64_t> input_shape_buffer(&kernel, input_shape.size());
|
||||
RocmKernel::RocmAsyncBuffer<int64_t> input_strides_buffer(&kernel, input_strides.size());
|
||||
for (int i = 0; i < input_shape.size(); i++) {
|
||||
input_shape_buffer.CpuPtr()[i] = input_shape[i];
|
||||
input_strides_buffer.CpuPtr()[i] = input_strides[i];
|
||||
}
|
||||
|
||||
ORT_RETURN_IF_ERROR(input_shape_buffer.CopyToGpu());
|
||||
ORT_RETURN_IF_ERROR(input_strides_buffer.CopyToGpu());
|
||||
|
||||
switch (element_size) {
|
||||
case sizeof(int8_t):
|
||||
hipLaunchKernelGGL(Transpose3DKernel<int8_t>, grid_size, block_size, 0, 0,
|
||||
input_shape_buffer.GpuPtr(), input_strides_buffer.GpuPtr(),
|
||||
reinterpret_cast<const ToHipType<int8_t>::MappedType*>(input_data),
|
||||
reinterpret_cast<ToHipType<int8_t>::MappedType*>(output_data));
|
||||
break;
|
||||
case sizeof(int16_t):
|
||||
hipLaunchKernelGGL(Transpose3DKernel<int16_t>, grid_size, block_size, 0, 0,
|
||||
input_shape_buffer.GpuPtr(), input_strides_buffer.GpuPtr(),
|
||||
reinterpret_cast<const ToHipType<int16_t>::MappedType*>(input_data),
|
||||
reinterpret_cast<ToHipType<int16_t>::MappedType*>(output_data));
|
||||
break;
|
||||
case sizeof(int32_t):
|
||||
hipLaunchKernelGGL(Transpose3DKernel<int32_t>, grid_size, block_size, 0, 0,
|
||||
input_shape_buffer.GpuPtr(), input_strides_buffer.GpuPtr(),
|
||||
reinterpret_cast<const ToHipType<int32_t>::MappedType*>(input_data),
|
||||
reinterpret_cast<ToHipType<int32_t>::MappedType*>(output_data));
|
||||
break;
|
||||
case sizeof(int64_t):
|
||||
hipLaunchKernelGGL(Transpose3DKernel<int64_t>, grid_size, block_size, 0, 0,
|
||||
input_shape_buffer.GpuPtr(), input_strides_buffer.GpuPtr(),
|
||||
reinterpret_cast<const ToHipType<int64_t>::MappedType*>(input_data),
|
||||
reinterpret_cast<ToHipType<int64_t>::MappedType*>(output_data));
|
||||
break;
|
||||
default:
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ",
|
||||
element_size);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
|
||||
template <int element_size>
|
||||
// __global__ void Transpose4DKernel(const TArray<int64_t> input_strides, const void* input_data,
|
||||
// const TArray<int64_t> output_strides, void* output_data,
|
||||
// HIP_LONG N) {
|
||||
__global__ void Transpose4DKernel(const int64_t* input_strides, const void* input_data,
|
||||
const int64_t* output_strides, void* output_data,
|
||||
HIP_LONG N) {
|
||||
// output coordinates will be: blockIdx.y, blockIdx.x, threadIdx.y, threadIdx.x
|
||||
HIP_LONG input_index = (blockIdx.y * input_strides[0] +
|
||||
blockIdx.x * input_strides[1] +
|
||||
threadIdx.y * input_strides[2]) /
|
||||
(4 * sizeof(int) / element_size) +
|
||||
threadIdx.x * input_strides[3];
|
||||
|
||||
HIP_LONG output_index = (blockIdx.y * output_strides[0] +
|
||||
blockIdx.x * output_strides[1] +
|
||||
threadIdx.y * output_strides[2]) /
|
||||
(4 * sizeof(int) / element_size) +
|
||||
threadIdx.x * output_strides[3];
|
||||
|
||||
const int4* v_input = reinterpret_cast<const int4*>(input_data);
|
||||
int4* v_output = reinterpret_cast<int4*>(output_data);
|
||||
|
||||
if (input_index < N && output_index < N) {
|
||||
v_output[output_index] = v_input[input_index];
|
||||
}
|
||||
}
|
||||
|
||||
bool CanDoTranspose4D(const hipDeviceProp_t& prop,
|
||||
size_t element_size,
|
||||
int32_t rank,
|
||||
const std::vector<int64_t>& input_dims,
|
||||
const std::vector<size_t>& permutations) {
|
||||
if (rank == 4 &&
|
||||
// the permutations is not on the last dimension.
|
||||
permutations[rank - 1] == (rank - 1)) {
|
||||
// The block size will be set based on the last two dimensions of 4D tensor.
|
||||
// the number threads per block will be calculated as below.
|
||||
int num_elements_per_thread = 4 * sizeof(int) / element_size; // int4 is used in the kernel to access data.
|
||||
int64_t num_elements_in_last_two_dimensions = input_dims[rank - 2] * input_dims[rank - 1];
|
||||
int64_t num_threads_per_block = num_elements_in_last_two_dimensions / num_elements_per_thread;
|
||||
|
||||
if (((num_elements_in_last_two_dimensions & (num_elements_per_thread - 1)) == 0) &&
|
||||
num_threads_per_block <= prop.maxThreadsPerBlock &&
|
||||
num_threads_per_block >= prop.warpSize &&
|
||||
// num_threads_per_block must be aligned with warp size: 32
|
||||
((num_threads_per_block & (prop.warpSize - 1)) == 0)) {
|
||||
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
// Status Transpose4DImpl(size_t element_size, const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides, const void* input_data,
|
||||
// const TArray<int64_t>& output_strides, void* output_data, int64_t N) {
|
||||
Status Transpose4DImpl(const Transpose& kernel, size_t element_size, const std::vector<int64_t>& input_shape, const std::vector<int64_t>& input_strides, const void* input_data,
|
||||
const std::vector<int64_t>& output_strides, void* output_data, int64_t N) {
|
||||
|
||||
int num_elements_per_thread = 4 * sizeof(int) / element_size; // int4 is used in the kernel to access data.
|
||||
dim3 block_size(input_shape[3] / num_elements_per_thread, input_shape[2]);
|
||||
dim3 grid_size(input_shape[1], input_shape[0]);
|
||||
|
||||
RocmKernel::RocmAsyncBuffer<int64_t> input_strides_buffer(&kernel, input_strides.size());
|
||||
RocmKernel::RocmAsyncBuffer<int64_t> output_strides_buffer(&kernel, output_strides.size());
|
||||
for (int i = 0; i < input_strides.size(); i++) {
|
||||
input_strides_buffer.CpuPtr()[i] = input_strides[i];
|
||||
output_strides_buffer.CpuPtr()[i] = output_strides[i];
|
||||
}
|
||||
|
||||
ORT_RETURN_IF_ERROR(input_strides_buffer.CopyToGpu());
|
||||
ORT_RETURN_IF_ERROR(output_strides_buffer.CopyToGpu());
|
||||
|
||||
switch (element_size) {
|
||||
case sizeof(int8_t):
|
||||
hipLaunchKernelGGL(Transpose4DKernel<sizeof(int8_t)>, grid_size, block_size, 0, 0,
|
||||
input_strides_buffer.GpuPtr(), input_data,
|
||||
output_strides_buffer.GpuPtr(), output_data, N / num_elements_per_thread);
|
||||
break;
|
||||
case sizeof(int16_t):
|
||||
hipLaunchKernelGGL(Transpose4DKernel<sizeof(int16_t)>, grid_size, block_size, 0, 0,
|
||||
input_strides_buffer.GpuPtr(), input_data,
|
||||
output_strides_buffer.GpuPtr(), output_data, N / num_elements_per_thread);
|
||||
break;
|
||||
case sizeof(int32_t):
|
||||
hipLaunchKernelGGL(Transpose4DKernel<sizeof(int32_t)>, grid_size, block_size, 0, 0,
|
||||
input_strides_buffer.GpuPtr(), input_data,
|
||||
output_strides_buffer.GpuPtr(), output_data, N / num_elements_per_thread);
|
||||
break;
|
||||
case sizeof(int64_t):
|
||||
hipLaunchKernelGGL(Transpose4DKernel<sizeof(int64_t)>, grid_size, block_size, 0, 0,
|
||||
input_strides_buffer.GpuPtr(), input_data,
|
||||
output_strides_buffer.GpuPtr(), output_data, N / num_elements_per_thread);
|
||||
break;
|
||||
default:
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for transpose on CUDA. Element size was ",
|
||||
element_size);
|
||||
}
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
// __global__ void _TransposeKernel(int32_t shape_rank, const TArray<int64_t> input_strides,
|
||||
// const T* __restrict__ input_data, const TArray<fast_divmod> output_strides, T* __restrict__ output_data, HIP_LONG N) {
|
||||
|
|
|
|||
|
|
@ -4,10 +4,27 @@
|
|||
#pragma once
|
||||
#include <stdint.h>
|
||||
#include "core/providers/rocm/shared_inc/rocm_utils.h"
|
||||
#include "core/providers/rocm/tensor/transpose.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace rocm {
|
||||
|
||||
bool CanDoTranspose3D(int32_t rank, const std::vector<int64_t>& input_dims, const std::vector<size_t>& permutations);
|
||||
// Status Transpose3DImpl(size_t element_size, const TArray<int64_t>& input_shape, const TArray<int64_t>& input_strides, const void* input_data,
|
||||
// void* output_data, int64_t N);
|
||||
Status Transpose3DImpl(const Transpose& kernel, size_t element_size, const std::vector<int64_t>& input_shape, const std::vector<int64_t>& input_strides, const void* input_data,
|
||||
void* output_data, int64_t N);
|
||||
|
||||
|
||||
bool CanDoTranspose4D(const hipDeviceProp_t& prop,
|
||||
size_t element_size,
|
||||
int32_t rank,
|
||||
const std::vector<int64_t>& input_dims,
|
||||
const std::vector<size_t>& permutations);
|
||||
|
||||
Status Transpose4DImpl(const Transpose& kernel, size_t element_size, const std::vector<int64_t>& input_shape, const std::vector<int64_t>& input_strides, const void* input_data,
|
||||
const std::vector<int64_t>& output_strides, void* output_data, int64_t N);
|
||||
|
||||
// Status TransposeImpl(size_t element_size, int32_t shape_rank, const TArray<int64_t>& input_strides,
|
||||
// const void* input_data, const TArray<fast_divmod>& fdm_output_strides, void* output_data, int64_t N);
|
||||
Status TransposeImpl(size_t element_size, int32_t shape_rank, const int64_t* input_strides,
|
||||
|
|
|
|||
Loading…
Reference in a new issue