1D, 2D and 3D transpose cuda
This commit is contained in:
lucasdelimanogueira
parent
fb6429aad6
commit
702ebaf87c
5 changed files with 123 additions and 17 deletions
|
|
@ -375,14 +375,14 @@ void transpose_2D_tensor_cpu(Tensor* tensor, float* result_data) {
|
|||
}
|
||||
|
||||
void transpose_3D_tensor_cpu(Tensor* tensor, float* result_data) {
|
||||
int depth = tensor->shape[0];
|
||||
int batch = tensor->shape[0];
|
||||
int rows = tensor->shape[1];
|
||||
int cols = tensor->shape[2];
|
||||
|
||||
for (int i = 0; i < depth; i++) {
|
||||
for (int i = 0; i < batch; i++) {
|
||||
for (int j = 0; j < rows; j++) {
|
||||
for (int k = 0; k < cols; k++) {
|
||||
result_data[k * rows * depth + j * depth + i] = tensor->data[i * rows * cols + j * cols + k];
|
||||
result_data[k * rows * batch + j * batch + i] = tensor->data[i * rows * cols + j * cols + k];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1011,23 +1011,46 @@ __host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data) {
|
|||
cudaDeviceSynchronize();
|
||||
}
|
||||
|
||||
__global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) {
|
||||
int tid_x = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int tid_y = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
|
||||
if (tid_x < cols && tid_y < rows) {
|
||||
result_data[tid_x * rows + tid_y] = data[tid_y * cols + tid_x];
|
||||
__global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int size) {
|
||||
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < size) {
|
||||
result_data[i] = data[i];
|
||||
}
|
||||
}
|
||||
|
||||
__host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
|
||||
__host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data) {
|
||||
|
||||
int number_of_blocks = (tensor->size + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK;
|
||||
transpose_1D_tensor_cuda_kernel<<<number_of_blocks, THREADS_PER_BLOCK>>>(tensor->data, result_data, tensor->size);
|
||||
|
||||
cudaError_t error = cudaGetLastError();
|
||||
if (error != cudaSuccess) {
|
||||
printf("CUDA error: %s\n", cudaGetErrorString(error));
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
cudaDeviceSynchronize();
|
||||
}
|
||||
|
||||
__global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int j = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
|
||||
if (i < rows && j < cols) {
|
||||
result_data[j * rows + i] = data[i * cols + j];
|
||||
}
|
||||
}
|
||||
|
||||
__host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data) {
|
||||
|
||||
int rows = tensor->shape[0];
|
||||
int cols = tensor->shape[1];
|
||||
|
||||
dim3 threadsPerBlock(16, 16);
|
||||
dim3 number_of_blocks((cols + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y);
|
||||
transpose_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
|
||||
dim3 number_of_blocks((rows + threadsPerBlock.x - 1) / threadsPerBlock.x, (cols + threadsPerBlock.y - 1) / threadsPerBlock.y);
|
||||
transpose_2D_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
|
||||
|
||||
|
||||
cudaError_t error = cudaGetLastError();
|
||||
|
|
@ -1039,6 +1062,37 @@ __host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data) {
|
|||
cudaDeviceSynchronize();
|
||||
}
|
||||
|
||||
__global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int batch, int rows, int cols) {
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int j = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
int k = blockIdx.z * blockDim.z + threadIdx.z;
|
||||
|
||||
if (i < batch && j < rows && k < cols) {
|
||||
result_data[k * rows * batch + j * batch + i] = data[i * rows * cols + j * cols + k];
|
||||
}
|
||||
}
|
||||
|
||||
__host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data) {
|
||||
|
||||
int batch = tensor->shape[0];
|
||||
int rows = tensor->shape[1];
|
||||
int cols = tensor->shape[2];
|
||||
|
||||
dim3 threadsPerBlock(8, 8, 8);
|
||||
dim3 number_of_blocks((batch + threadsPerBlock.x - 1) / threadsPerBlock.x, (rows + threadsPerBlock.y - 1) / threadsPerBlock.y, (cols + threadsPerBlock.z - 1) / threadsPerBlock.z);
|
||||
transpose_2D_tensor_cuda_kernel<<<number_of_blocks, threadsPerBlock>>>(tensor->data, result_data, rows, cols);
|
||||
|
||||
|
||||
cudaError_t error = cudaGetLastError();
|
||||
if (error != cudaSuccess) {
|
||||
printf("CUDA error: %s\n", cudaGetErrorString(error));
|
||||
exit(-1);
|
||||
}
|
||||
|
||||
cudaDeviceSynchronize();
|
||||
}
|
||||
|
||||
|
||||
__global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size) {
|
||||
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
|
|
|||
|
|
@ -73,8 +73,14 @@
|
|||
__global__ void zeros_like_tensor_cuda_kernel(float* data, float* result_data, int size);
|
||||
__host__ void zeros_like_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
|
||||
__global__ void transpose_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
|
||||
__host__ void transpose_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
__global__ void transpose_1D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
|
||||
__host__ void transpose_1D_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
|
||||
__global__ void transpose_2D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
|
||||
__host__ void transpose_2D_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
|
||||
__global__ void transpose_3D_tensor_cuda_kernel(float* data, float* result_data, int rows, int cols);
|
||||
__host__ void transpose_3D_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
|
||||
__global__ void assign_tensor_cuda_kernel(float* data, float* result_data, int size);
|
||||
__host__ void assign_tensor_cuda(Tensor* tensor, float* result_data);
|
||||
|
|
|
|||
|
|
@ -1422,7 +1422,20 @@ extern "C" {
|
|||
|
||||
float* result_data;
|
||||
cudaMalloc((void **)&result_data, size * sizeof(float));
|
||||
transpose_tensor_cuda(tensor, result_data);
|
||||
switch (ndim) {
|
||||
case 1:
|
||||
transpose_1D_tensor_cuda(tensor, result_data);
|
||||
break;
|
||||
case 2:
|
||||
transpose_2D_tensor_cuda(tensor, result_data);
|
||||
break;
|
||||
case 3:
|
||||
transpose_3D_tensor_cuda(tensor, result_data);
|
||||
break;
|
||||
default:
|
||||
fprintf(stderr, "Transpose only supports tensors up to 3 dimensions.\n");
|
||||
exit(-1);
|
||||
}
|
||||
return create_tensor(result_data, shape, ndim, device);
|
||||
}
|
||||
else {
|
||||
|
|
@ -1479,7 +1492,15 @@ extern "C" {
|
|||
float* result_data;
|
||||
cudaMalloc((void **)&result_data, size * sizeof(float));
|
||||
assign_tensor_cuda(tensor, result_data);
|
||||
return create_tensor(result_data, shape, ndim, device);
|
||||
|
||||
Tensor* new_tensor = create_tensor(result_data, shape, ndim, device);
|
||||
for (int i = 0; i < ndim; i++) {
|
||||
new_tensor->strides[i] = tensor->strides[i];
|
||||
}
|
||||
new_tensor->strides[axis1] = tensor->strides[axis2];
|
||||
new_tensor->strides[axis2] = tensor->strides[axis1];
|
||||
make_contiguous(new_tensor);
|
||||
return new_tensor;
|
||||
}
|
||||
else {
|
||||
float* result_data = (float*)malloc(size * sizeof(float));
|
||||
|
|
|
|||
|
|
@ -471,8 +471,33 @@ class TestTensorOperations(unittest.TestCase):
|
|||
|
||||
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
|
||||
|
||||
def test_1D_T(self):
|
||||
"""
|
||||
Test transposition of a 1D tensor.
|
||||
"""
|
||||
norch_tensor = norch.Tensor([1, 2, 3, 4]).to(self.device)
|
||||
norch_result = norch_tensor.T
|
||||
torch_result = utils.to_torch(norch_result).to(self.device)
|
||||
|
||||
def test_transpose_T(self):
|
||||
torch_tensor = torch.tensor([1, 2, 3, 4]).to(self.device)
|
||||
torch_expected = torch_tensor.T
|
||||
|
||||
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
|
||||
|
||||
def test_2D_T(self):
|
||||
"""
|
||||
Test transposition of a 2D tensor.
|
||||
"""
|
||||
norch_tensor = norch.Tensor([[1, 2, 3], [4, 5, 6]]).to(self.device)
|
||||
norch_result = norch_tensor.T
|
||||
torch_result = utils.to_torch(norch_result).to(self.device)
|
||||
|
||||
torch_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]).to(self.device)
|
||||
torch_expected = torch_tensor.T
|
||||
|
||||
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
|
||||
|
||||
def test_3D_T(self):
|
||||
"""
|
||||
Test transposition of a tensor: tensor.T
|
||||
"""
|
||||
|
|
@ -481,7 +506,7 @@ class TestTensorOperations(unittest.TestCase):
|
|||
torch_result = utils.to_torch(norch_result).to(self.device)
|
||||
|
||||
torch_tensor = torch.tensor([[[1, 2], [3, -4]], [[5, 6], [7, 8]]]).to(self.device)
|
||||
torch_expected = torch.transpose(torch_tensor, 0, 2)
|
||||
torch_expected = torch_tensor.T
|
||||
|
||||
self.assertTrue(utils.compare_torch(torch_result, torch_expected))
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue