optimize embedding cuda kernel lookup_table_v2,test=develop (#25587)

fix_copy_if_different
donproc 5 years ago committed by GitHub
parent 7206417259
commit 95fa383df2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -105,17 +105,17 @@ class LookupTableV2CUDAKernel : public framework::OpKernel<T> {
auto *table = table_t->data<T>();
auto *output = output_t->mutable_data<T>(context.GetPlace());
dim3 threads(128, 8);
dim3 grids(8, 1);
dim3 threads(256, 4);
dim3 grids(80, 1);
if (padding_idx == -1)
LookupTableV2<
T, 128, 8, 8,
T, 256, 4, 80,
false><<<grids, threads, 0, context.cuda_device_context().stream()>>>(
output, table, ids, N, K, D, padding_idx);
else
LookupTableV2<
T, 128, 8, 8,
T, 256, 4, 80,
true><<<grids, threads, 0, context.cuda_device_context().stream()>>>(
output, table, ids, N, K, D, padding_idx);
}

Loading…
Cancel
Save