saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Update GratherGard to accumulate in fp32 (#4601)

Browse source
This commit is contained in:
Sherlock 2020-07-24 10:54:31 -07:00 committed by GitHub
parent 9c75c29403
commit aa328c2c20
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
orttraining/orttraining/training_ops/cuda/tensor/gather_grad_impl.cu
View file

@ -40,15 +40,15 @@ __global__ void _GatherGradImpl(
const int weight_row = itr * input_numel + ((int)input[idx]) * stride; //the offset of the input
const int grad_row = (itr * numel + ((int)indices[idx])) * stride; //the offset of the gradient
T gradient[SZ];
T weight[SZ];
float gradient[SZ];
float weight[SZ];
#pragma unroll
for (int ii = 0; ii < SZ; ii++) {
int feature_dim = start_feature + ii * GPU_WARP_SIZE;
if (feature_dim < stride) {
gradient[ii] = static_cast<T>(grad_output[grad_row + feature_dim]);
weight[ii] = static_cast<T>(grad_weight[weight_row + feature_dim]);
gradient[ii] = static_cast<float>(grad_output[grad_row + feature_dim]);
weight[ii] = static_cast<float>(grad_weight[weight_row + feature_dim]);
}
}