Fix GridEncoding gradient with max_level / Interpolation::Nearest

YOUSIKI · Feb 11, 2022 · d829746 · d829746
1 parent 9b9d046
commit d829746
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/include/tiny-cuda-nn/encodings/grid.h b/include/tiny-cuda-nn/encodings/grid.h
@@ -151,6 +151,15 @@ __global__ void kernel_grid(
 			encoded_positions[i + (level * N_FEATURES_PER_LEVEL + f) * num_elements] = (T)0.0f;
 		}
 
+		// Gradient is zero for zeroed-out dimensions.
+		if (dy_dx) {
+			#pragma unroll
+			for (uint32_t grad_dim = 0; grad_dim < N_POS_DIMS; ++grad_dim) {
+				const uint32_t fan_out_grad = num_grid_features * N_POS_DIMS;
+				*(vector_fullp_t<N_FEATURES_PER_LEVEL>*)&dy_dx[i * fan_out_grad + level * N_FEATURES_PER_LEVEL + grad_dim * num_grid_features] = {0};
+			}
+		}
+
 		return;
 	}
 
@@ -189,6 +198,15 @@ __global__ void kernel_grid(
 			encoded_positions[i + (level * N_FEATURES_PER_LEVEL + f) * num_elements] = result[f];
 		}
 
+		// Gradient is zero when there's no interpolation.
+		if (dy_dx) {
+			#pragma unroll
+			for (uint32_t grad_dim = 0; grad_dim < N_POS_DIMS; ++grad_dim) {
+				const uint32_t fan_out_grad = num_grid_features * N_POS_DIMS;
+				*(vector_fullp_t<N_FEATURES_PER_LEVEL>*)&dy_dx[i * fan_out_grad + level * N_FEATURES_PER_LEVEL + grad_dim * num_grid_features] = {0};
+			}
+		}
+
 		return;
 	}