[Prim] Support masked_fill_double_grad for eager mode (#73140)

HydrogenSulfate · web-flow · commit 533f1a0a9cbb · 2025-06-10T19:55:30.000+08:00
* support masked_fill_double_grad for eager mode

* update code

* remove redundant code

* Update op_compat.yaml
diff --git a/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py b/paddle/fluid/eager/auto_code_generator/generator/eager_gen.py
@@ -91,6 +91,7 @@
     "index_add_double_grad",
     "acos_double_grad",
     "put_along_axis_double_grad",
+    "masked_fill_double_grad",
 ]
 
 # white ops list whose kernel can automatically do type promotion.
diff --git a/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py b/paddle/fluid/pir/dialect/op_generator/vjp_interface_black_list.py
@@ -39,4 +39,5 @@
     'index_add_grad',
     'acos_grad',
     'put_along_axis_grad',
+    'masked_fill_grad',
 ]
diff --git a/paddle/fluid/prim/api/api.yaml b/paddle/fluid/prim/api/api.yaml
@@ -57,3 +57,4 @@
 - isfinite
 - take_along_axis
 - variance
+- masked_fill
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -2085,60 +2085,6 @@ void pad_grad(const Tensor& input,
   }
 }
 
-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) {
-  auto x_grad_dims = x.dims();
-  auto mask_dims = mask.dims();
-  auto expanded_size = common::vectorize(
-      phi::funcs::BroadcastTwoDims(x_grad_dims, mask_dims, -1));
-  auto expanded_dims = common::make_ddim(expanded_size);
-  auto mask_expand = Tensor();
-
-  if (mask_dims != x_grad_dims) {
-    mask_expand = expand<T>(mask, expanded_size);
-  } else {
-    mask_expand = mask;
-  }
-
-  if (x_grad) {
-    bool expand_x = false;
-    auto x_grad_tmp = Tensor();
-    auto x_grad_expand = Tensor();
-    auto value_expand = Tensor();
-
-    if (x_grad->dims() != expanded_dims) {
-      x_grad_expand = full<T>(expanded_size, 0.0, x.dtype(), x.place());
-      expand_x = true;
-    } else {
-      x_grad_expand = *x_grad;
-    }
-
-    value_expand = full<T>(expanded_size, 0.0, x.dtype(), x.place());
-
-    x_grad_tmp = *x_grad;
-    if (expand_x) {
-      x_grad_tmp = x_grad_expand;
-    }
-
-    x_grad_tmp = where<T>(mask_expand, value_expand, out_grad);
-
-    if (expand_x) {
-      expand_grad<T>(x, x_grad_expand, IntArray(expanded_size), &x_grad_tmp);
-    }
-    set_output<T>(x_grad_tmp, x_grad);
-  }
-
-  if (value_grad) {
-    auto value_grad_tmp = sum<T>(mask_expand);
-    set_output<T>(value_grad_tmp, value_grad);
-  }
-}
-
 template <typename T>
 void scatter_nd_add_grad(const Tensor& index,
                          const Tensor& updates,
diff --git a/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_double_backward_api.h
@@ -156,6 +156,29 @@ void pow_double_grad(const Tensor& x,
   }
 }
 
+template <typename T>
+void masked_fill_double_grad(const Tensor& mask,
+                             const paddle::optional<Tensor>& grad_x_grad,
+                             const paddle::optional<Tensor>& grad_value_grad,
+                             Tensor* grad_out_grad) {
+  if (grad_out_grad) {
+    Tensor grad_out_grad_tmp;
+    if (grad_x_grad && grad_value_grad) {
+      grad_out_grad_tmp =
+          masked_fill<T>(grad_x_grad.get(), mask, grad_value_grad.get());
+    } else if (grad_x_grad) {
+      grad_out_grad_tmp = masked_fill<T>(
+          grad_x_grad.get(),
+          mask,
+          full<T>({}, 0, grad_x_grad.get().dtype(), grad_x_grad.get().place()));
+    } else if (grad_value_grad) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "grad_x_grad can not be null in 'masked_fill_double_grad'"));
+    }
+    set_output<T>(grad_out_grad_tmp, grad_out_grad);
+  }
+}
+
 template <typename T>
 void maximum_double_grad(const Tensor& x,
                          const Tensor& y,
diff --git a/paddle/fluid/primitive/codegen/decomp_vjp_gen.py b/paddle/fluid/primitive/codegen/decomp_vjp_gen.py
@@ -128,7 +128,6 @@
     'subtract_grad',
     'sum_grad',
     'take_along_axis_grad',
-    'put_along_axis_grad',
     'tanh_grad',
     'tile_grad',
     'topk_grad',
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2134,6 +2134,16 @@
     data_type : softmax
   inplace : (softmax -> logits_grad)
 
+- backward_op : masked_fill_double_grad
+  forward : masked_fill_grad (Tensor x, Tensor mask, Tensor value, Tensor grad_out) -> Tensor(grad_x), Tensor(grad_value)
+  args : (Tensor mask, Tensor grad_x_grad, Tensor grad_value_grad)
+  output : Tensor(grad_out_grad)
+  infer_meta :
+    func : UnchangedInferMeta
+    param : [grad_x_grad]
+  optional: grad_x_grad, grad_value_grad
+  composite : masked_fill_double_grad(mask, grad_x_grad, grad_value_grad, grad_out_grad)
+
 - backward_op : masked_fill_grad
   forward : masked_fill (Tensor x, Tensor mask, Tensor value) -> Tensor(out)
   args : (Tensor x, Tensor mask, Tensor value, Tensor out_grad)
@@ -2145,7 +2155,7 @@
     func : masked_fill_grad
     param: [x, mask, value, out_grad]
   inplace : (out_grad -> x_grad)
-  composite : masked_fill_grad(x, mask, value, out_grad, x_grad, value_grad)
+  backward: masked_fill_double_grad
 
 - backward_op : masked_select_grad
   forward : masked_select (Tensor x, Tensor mask) -> Tensor(out)

Original file line number	Diff line number	Diff line change
`@@ -91,6 +91,7 @@`
`91`	`91`	`"index_add_double_grad",`
`92`	`92`	`"acos_double_grad",`
`93`	`93`	`"put_along_axis_double_grad",`
	`94`	`+ "masked_fill_double_grad",`
`94`	`95`	`]`
`95`	`96`
`96`	`97`	`# white ops list whose kernel can automatically do type promotion.`
Original file line number	Diff line number	Diff line change
`@@ -39,4 +39,5 @@`
`39`	`39`	`'index_add_grad',`
`40`	`40`	`'acos_grad',`
`41`	`41`	`'put_along_axis_grad',`
	`42`	`+ 'masked_fill_grad',`
`42`	`43`	`]`