[Sparse] Fix bugs in parameter freezing (#56154)

* Add enforce for sparse_bn.

* Add enforce for sp conv.

paddle/phi/kernels/funcs/sparse/convolution.h

@@ -163,6 +163,8 @@ inline void SubmPreProcess(const Context& dev_ctx,
                            DenseTensor* kernel_grad,
                            DenseTensor* x_grad) {
   auto blas = phi::funcs::GetBlas<Context, T>(dev_ctx);
+  const bool is_params_freezing = kernel_grad == nullptr;
+  if (!is_params_freezing) {
     T* d_kernel_ptr = kernel_grad->data<T>();
     blas.GEMM(CblasTrans,
               CblasNoTrans,
@@ -174,6 +176,7 @@ inline void SubmPreProcess(const Context& dev_ctx,
               out_grad.data<T>(),
               static_cast<T>(0),
               d_kernel_ptr + half_kernel_size * in_channels * out_channels);
+  }
 
   // call gemm: d_x = out_grad * transpose(kernel)
   // (n, out_channels) * (out_channels, in_channels)

paddle/phi/kernels/sparse/batch_norm_grad_kernel.cc

@@ -42,8 +42,19 @@ void BatchNormCooGradKernel(const Context& dev_ctx,
                             DenseTensor* scale_grad,
                             DenseTensor* bias_grad) {
   EmptyLikeCooKernel<T, Context>(dev_ctx, x, x_grad);
+
+  // TODO(umiswing): add check for parameter freezing automatically
+  PADDLE_ENFORCE_EQ((scale_grad == nullptr && bias_grad == nullptr) ||
+                        (scale_grad != nullptr && bias_grad != nullptr),
+                    true,
+                    phi::errors::InvalidArgument(
+                        "Weight and bias's stop_gradient of BatchNorm must be "
+                        "True or False at the same time."));
+
+  if (scale_grad && bias_grad) {
     *scale_grad = phi::EmptyLike<T, Context>(dev_ctx, scale);
     *bias_grad = phi::EmptyLike<T, Context>(dev_ctx, bias);
+  }
   phi::BatchNormGradKernel<T, Context>(dev_ctx,
                                        x.values(),
                                        scale,

paddle/phi/kernels/sparse/gpu/conv_grad_kernel.cu

@@ -56,6 +56,7 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
                             const std::string& key,
                             SparseCooTensor* x_grad,
                             DenseTensor* kernel_grad) {
+  const bool is_params_freezing = kernel_grad == nullptr;
   const auto& kernel_dims = kernel.dims();
   const bool is2D = kernel_dims.size() == 4 ? true : false;
   const int kernel_size =
@@ -79,10 +80,13 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
   T* in_features_ptr = in_features.data<T>();
   T* d_x_features_ptr = d_x_features.data<T>();
   T* out_grad_features_ptr = out_grad_features.data<T>();
+  T* d_kernel_ptr = nullptr;
+  if (!is_params_freezing) {
     *kernel_grad = phi::EmptyLike<T>(dev_ctx, kernel);
-  T* d_kernel_ptr = kernel_grad->data<T>();
+    d_kernel_ptr = kernel_grad->data<T>();
     phi::backends::gpu::GpuMemsetAsync(
         d_kernel_ptr, 0, sizeof(T) * kernel_grad->numel(), dev_ctx.stream());
+  }
 
   int half_kernel_size = kernel_size / 2;
   auto blas = phi::funcs::GetBlas<GPUContext, T>(dev_ctx);
@@ -184,6 +188,8 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
   }
 #endif
   const T* kernel_ptr = kernel.data<T>();
+  T* tmp_d_x_ptr = nullptr;
+  T* tmp_d_kernel_ptr = nullptr;
   for (int i = 0; i < kernel_size; i++) {
     if (counter_ptr[i] <= 0 || (subm && i == half_kernel_size)) {
       continue;
@@ -195,8 +201,10 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
     T* tmp_in_ptr = in_features_ptr + offsets[i] * in_channels;
     T* tmp_out_grad_ptr = out_grad_features_ptr + offsets[i] * out_channels;
     const T* tmp_kernel_ptr = kernel_ptr + i * in_channels * out_channels;
-    T* tmp_d_x_ptr = d_x_features_ptr + offsets[i] * in_channels;
-    T* tmp_d_kernel_ptr = d_kernel_ptr + i * in_channels * out_channels;
+    tmp_d_x_ptr = d_x_features_ptr + offsets[i] * in_channels;
+    if (!is_params_freezing) {
+      tmp_d_kernel_ptr = d_kernel_ptr + i * in_channels * out_channels;
+    }
 
 #if defined(PADDLE_WITH_CUTLASS) && SPCONV_WITH_CUTLASS
     if (cutlass) {
@@ -204,6 +212,7 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
       const IntT* scatter_x_indices = rulebook_ptr + offsets[i];
       const IntT* gather_out_indices = rulebook_ptr + rulebook_len + offsets[i];
       const size_t key = autotune::GenKey(M / features_num_range, N, K);
+      if (!is_params_freezing) {
         // call gemm: d_kernel = transpose(x) * out_grad
         // (in_channels, n) * (n, out_channels)
         static cutlass::device_memory::allocation<uint8_t> workspace(
@@ -224,6 +233,7 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
             static_cast<const T>(1.0),
             static_cast<const T>(0.0),
             &workspace);
+      }
       // call gemm: d_x = out_grad * transpose(kernel)
       // (n, out_channels) * (out_channels, in_channels)
       GatherGemmScatterDriver<80, false, true>(
@@ -244,6 +254,7 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
           nullptr);
     } else {
 #endif
+      if (!is_params_freezing) {
         // call gemm: d_kernel = transpose(x) * out_grad
         // (in_channels, n) * (n, out_channels)
         blas.GEMM(CblasTrans,
@@ -256,6 +267,7 @@ void Conv3dCooGradGPUKernel(const GPUContext& dev_ctx,
                   tmp_out_grad_ptr,
                   static_cast<T>(0),
                   tmp_d_kernel_ptr);
+      }
 
       // call gemm: d_x = out_grad * transpose(kernel)
       // (n, out_channels) * (out_channels, in_channels)