Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[backend](cuda): faster uncontiguous concat #10760

Merged
merged 4 commits into from
Dec 12, 2024
Merged

[backend](cuda): faster uncontiguous concat #10760

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e4189e3
faster uncontiguous concat
A3shTnT Dec 10, 2024
f8a5b04
Use a lambda to avoid code duplication
A3shTnT Dec 11, 2024
618708c
Update ggml/src/ggml-cuda/concat.cu
slaren Dec 11, 2024
9fdf8ad
add constexpr and static assert
A3shTnT Dec 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 38 additions & 14 deletions ggml/src/ggml-cuda/concat.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,9 @@ static void concat_f32_cuda(const float * x, const float * y, float * dst, int n
}

// non-contiguous kernel (slow)
static __global__ void concat_f32_non_cont(
template <int dim>
static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE)
concat_f32_non_cont(
const char * src0,
const char * src1,
char * dst,
Expand All @@ -121,22 +123,26 @@ static __global__ void concat_f32_non_cont(
uint64_t nb0,
uint64_t nb1,
uint64_t nb2,
uint64_t nb3,
int32_t dim) {
uint64_t nb3){
const int64_t i3 = blockIdx.z;
const int64_t i2 = blockIdx.y;
const int64_t i1 = blockIdx.x;

int64_t o[4] = {0, 0, 0, 0};
o[dim] = dim == 0 ? ne00 : (dim == 1 ? ne01 : (dim == 2 ? ne02 : ne03));

const float * x;

for (int i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) {
for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) {
if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) {
x = (const float *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00);
} else {
x = (const float *)(src1 + (i3 - o[3])*nb13 + (i2 - o[2])*nb12 + (i1 - o[1])*nb11 + (i0 - o[0])*nb10);
if /*constexpr*/ (dim == 0) {
Copy link
Collaborator

@JohannesGaessler JohannesGaessler Dec 10, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a static_assert for dim.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also note that if constexpr should be available now that we are using C++17, no need to keep it commented.

Copy link
Contributor Author

@A3shTnT A3shTnT Dec 11, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also note that if constexpr should be available now that we are using C++17, no need to keep it commented.

Thank you for your advice, I removed the comments. I also found that the number and content of instructions generated by the compiler did not change with or without constexpr. Perhaps nvcc has already optimized for this point

x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + i1 * nb11 + (i0 - ne00) * nb10);
} else if (dim == 1) {
x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + (i1 - ne01) * nb11 + i0 * nb10);
} else if (dim == 2) {
x = (const float *) (src1 + i3 * nb13 + (i2 - ne02) * nb12 + i1 * nb11 + i0 * nb10);
} else if (dim == 3) {
x = (const float *) (src1 + (i3 - ne03) * nb13 + i2 * nb12 + i1 * nb11 + i0 * nb10);
}
}

float * y = (float *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0);
Expand Down Expand Up @@ -182,15 +188,33 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
}
} else {
dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]);
concat_f32_non_cont<<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
(const char *)src0->data,
(const char *)src1->data,
( char *)dst->data,
auto launch_kernel = [&](auto dim) {
concat_f32_non_cont<dim><<<grid_dim, CUDA_CONCAT_BLOCK_SIZE, 0, stream>>>(
(const char *) src0->data, (const char *) src1->data, (char *) dst->data,
src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3],
src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3],
src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3],
src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3],
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3], dim);
dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3],
dst->nb[0], dst->nb[1], dst->nb[2], dst->nb[3]);
};
switch (dim) {
case 0:
launch_kernel(std::integral_constant<int, 0>{});
break;
case 1:
launch_kernel(std::integral_constant<int, 1>{});
break;
case 2:
launch_kernel(std::integral_constant<int, 2>{});
break;
case 3:
launch_kernel(std::integral_constant<int, 3>{});
break;
default:
GGML_ABORT("Invalid dim: %d", dim);
break;
}
}
slaren marked this conversation as resolved.
Show resolved Hide resolved
}
}
Loading