[Comp] Support masked_fill_grad in comp (#73185)

HydrogenSulfate · web-flow · commit f78c9f1245fa · 2025-06-16T16:22:57.000+08:00
* support masked_fill_grad in comp

* fix

* restore

* update masked_fill_grad with where

* update code

* update code

* fix

* add no_need_buffer for masked_fill_grad
diff --git a/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py b/paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py
@@ -137,6 +137,7 @@
     'transpose_grad',
     'unsqueeze_grad',
     'p_norm_grad',
+    'masked_fill_grad',
 ]
 
 # declare belongs to codegen, but implementation not
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -161,6 +161,7 @@
     'elu_grad',
     'swiglu_grad',
     'p_norm_grad',
+    'masked_fill_grad',
 ]  # custom vjp list of composite op
 
 VJP_COMPS = PRIM_VJP + CUSTOM_VJP
diff --git a/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h b/paddle/fluid/primitive/decomp_rule/decomp_vjp/details.h
@@ -1075,6 +1075,78 @@ void softmax_grad(const Tensor& out,
   }
 }
 
+template <typename T>
+void masked_fill_grad(const Tensor& x,
+                      const Tensor& mask,
+                      const Tensor& value,
+                      const Tensor& out_grad,
+                      Tensor* x_grad,
+                      Tensor* value_grad) {
+  /**
+   * dx_i = dy_i if mask_i = 0 else 0
+   * dv = \sum{dy_i if mask_i = 1}
+   */
+  if (has_dynamic_shape(x.shape()) || has_dynamic_shape(mask.shape())) {
+    Tensor out_dims = shape64<T>(out_grad);
+    Tensor mask_expand = backend::expand<T>(mask, out_dims);
+    Tensor zeros_expand = backend::full_with_tensor<T>(
+        out_dims, 0, out_grad.dtype(), out_grad.place());
+
+    if (x_grad) {
+      Tensor x_grad_expand = where<T>(mask_expand, zeros_expand, out_grad);
+      Tensor x_grad_tmp = reduce_as<T>(x_grad_expand, x);
+      set_output<T>(x_grad_tmp, x_grad);
+    }
+
+    if (value_grad) {
+      Tensor value_grad_expand = where<T>(mask_expand, out_grad, zeros_expand);
+      Tensor value_grad_tmp =
+          sum<T>(value_grad_expand, {}, value.dtype(), false);
+      set_output<T>(value_grad_tmp, value_grad);
+    }
+
+  } else {
+    // expand mask to match x.shape
+    auto expanded_dims = out_grad.dims();
+    auto expanded_dims_vec = common::vectorize(expanded_dims);
+    Tensor mask_expand;
+    if (mask.dims() != expanded_dims) {
+      mask_expand = expand<T>(mask, expanded_dims_vec);
+    } else {
+      mask_expand = mask;
+    }
+    Tensor zeros_expand =
+        full<T>(expanded_dims_vec, 0, out_grad.dtype(), out_grad.place());
+
+    if (x_grad) {
+      Tensor x_grad_expand = where<T>(mask_expand, zeros_expand, out_grad);
+      if (x.dims() != x_grad_expand.dims()) {
+        auto reduce_dim =
+            get_reduce_dims_from_out(x_grad_expand.dims(), x.dims());
+        if (!reduce_dim.size()) {
+          set_output<T>(x_grad_expand, x_grad);
+        } else {
+          Tensor x_grad_tmp = sum<T>(
+              x_grad_expand, common::vectorize(reduce_dim), x.dtype(), false);
+          if (x_grad_tmp.dims() != x.dims()) {
+            x_grad_tmp = reshape<T>(x_grad_tmp, x.shape());
+          }
+          set_output<T>(x_grad_tmp, x_grad);
+        }
+      } else {
+        set_output<T>(x_grad_expand, x_grad);
+      }
+    }
+
+    if (value_grad) {
+      Tensor value_grad_expand = where<T>(mask_expand, out_grad, zeros_expand);
+      Tensor value_grad_tmp =
+          sum<T>(value_grad_expand, {}, value.dtype(), false);
+      set_output<T>(value_grad_tmp, value_grad);
+    }
+  }
+}
+
 template <typename T>
 void squeeze_grad(const Tensor& x,
                   const Tensor& out_grad,
diff --git a/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_grad_kernel.cu
@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #include "paddle/phi/kernels/masked_fill_grad_kernel.h"
+#include "paddle/phi/kernels/cast_kernel.h"
 #include "paddle/phi/kernels/funcs/masked_fill_utils.h"
 
 #include "paddle/phi/backends/gpu/gpu_context.h"
@@ -174,7 +175,12 @@ void MaskedFillGradKernel(const Context& dev_ctx,
       std::vector<int> mask_dims(mask.dims().size());
       std::iota(mask_dims.begin(), mask_dims.end(), 0);
       IntArray mask_axis(mask_dims);
-      SumKernel<T>(dev_ctx, mask, mask_axis, v_grad->dtype(), false, v_grad);
+      SumKernel<T>(dev_ctx,
+                   Cast<bool>(dev_ctx, mask, value.dtype()),
+                   mask_axis,
+                   v_grad->dtype(),
+                   false,
+                   v_grad);
       ScaleKernel<T>(dev_ctx,
                      *v_grad,
                      (out_grad.numel() / mask.numel()),
@@ -221,7 +227,12 @@ void MaskedFillGradKernel(const Context& dev_ctx,
     std::vector<int> v_dims(mask_expand.dims().size());
     std::iota(v_dims.begin(), v_dims.end(), 0);
     IntArray v_axis(v_dims);
-    SumKernel<T>(dev_ctx, mask_expand, v_axis, v_grad->dtype(), false, v_grad);
+    SumKernel<T>(dev_ctx,
+                 Cast<bool>(dev_ctx, mask_expand, value.dtype()),
+                 v_axis,
+                 v_grad->dtype(),
+                 false,
+                 v_grad);
   }
 }
 
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2160,6 +2160,7 @@
     func : masked_fill_grad
     param: [x, mask, value, out_grad]
   inplace : (out_grad -> x_grad)
+  no_need_buffer : x, value
   backward: masked_fill_double_grad
 
 - backward_op : masked_select_grad
diff --git a/python/paddle/autograd/backward_utils.py b/python/paddle/autograd/backward_utils.py
@@ -108,6 +108,7 @@
     "pd_op.where",
     "pd_op.p_norm",
     "pd_op.elu",
+    "pd_op.masked_fill",
 ]
 
 

Original file line number	Diff line number	Diff line change
`@@ -137,6 +137,7 @@`
`137`	`137`	`'transpose_grad',`
`138`	`138`	`'unsqueeze_grad',`
`139`	`139`	`'p_norm_grad',`
	`140`	`+ 'masked_fill_grad',`
`140`	`141`	`]`
`141`	`142`
`142`	`143`	`# declare belongs to codegen, but implementation not`
Original file line number	Diff line number	Diff line change
`@@ -108,6 +108,7 @@`
`108`	`108`	`"pd_op.where",`
`109`	`109`	`"pd_op.p_norm",`
`110`	`110`	`"pd_op.elu",`
	`111`	`+ "pd_op.masked_fill",`
`111`	`112`	`]`
`112`	`113`
`113`	`114`