[Prim][Decomp] masked_fill bug fix (#72819)

Eddie-Wang1120 · web-flow · commit 6283b96b511d · 2025-05-26T16:32:47.000+08:00
* add masked_fill test

* add masked_fill decomp

* fix cuda error

* fix bug

* fix bug

* fix inplace

* test benchmark

* add composite masked_fill_grad
diff --git a/paddle/fluid/prim/api/composite_backward/composite_backward_api.h b/paddle/fluid/prim/api/composite_backward/composite_backward_api.h
@@ -26,6 +26,7 @@
 #include "paddle/fluid/prim/api/generated_prim/prim_generated_api.h"
 #include "paddle/phi/common/amp_type_traits.h"
 #include "paddle/phi/common/int_array.h"
+#include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
 
 namespace paddle {
 namespace prim {
@@ -2084,6 +2085,60 @@ 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/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1925,8 +1925,7 @@ static PyObject* tensor__setitem_dygraph(TensorObject* self,
             mesh, self->tensor, transed_sub_tensor, value_tensor);
       }
       paddle::Tensor mask_tensor;
-      if (transed_index.size() == 1 &&
-          MaskedFillDispatching(
+      if (MaskedFillDispatching(
               transed_sub_tensor, value_tensor, transed_index, &mask_tensor)) {
         transed_sub_tensor =
             masked_fill__ad_func(transed_sub_tensor, mask_tensor, value_tensor);
diff --git a/paddle/fluid/pybind/slice_utils.h b/paddle/fluid/pybind/slice_utils.h
@@ -410,7 +410,7 @@ inline static bool MaskedFillDispatching(
     const paddle::Tensor& value,
     const std::vector<paddle::Tensor>& indices,
     paddle::Tensor* mask_tensor) {
-  if (value.numel() != 1) return false;
+  if (indices.size() != 1 || value.numel() != 1) return false;
   int64_t num_ind = 0;
   if ((indices)[0].dtype() != phi::DataType::BOOL) {
     return false;
diff --git a/paddle/phi/infermeta/binary.cc b/paddle/phi/infermeta/binary.cc
@@ -28,6 +28,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/unary.h"
 #include "paddle/phi/kernels/cpu/conv_util.h"
 #include "paddle/phi/kernels/funcs/axis_utils.h"
+#include "paddle/phi/kernels/funcs/common_infer_shape_functions.h"
 #include "paddle/phi/kernels/funcs/common_shape.h"
 #include "paddle/phi/kernels/funcs/correlation_funcs.h"
 
@@ -2906,7 +2907,10 @@ void MaskedFillInferMeta(const MetaTensor& x,
                          const MetaTensor& mask,
                          const MetaTensor& value,
                          MetaTensor* out) {
-  out->set_dims(x.dims());  // can not infer
+  auto x_dims = x.dims();
+  auto mask_dims = mask.dims();
+  auto expanded_dims = phi::funcs::BroadcastTwoDims(x_dims, mask_dims, -1);
+  out->set_dims(expanded_dims);
   out->set_dtype(x.dtype());
 }
 
diff --git a/paddle/phi/kernels/gpu/masked_fill_kernel.cu b/paddle/phi/kernels/gpu/masked_fill_kernel.cu
@@ -38,6 +38,11 @@ __global__ void GPUMaskedFillKernel(const T* input,
                                     const int64_t batch_size,
                                     T* output) {
   int64_t idx = (blockIdx.x * blockDim.x + threadIdx.x);
+
+  if (idx >= (input_len / VecSize)) {
+    return;
+  }
+
   int64_t vec_idx = idx * VecSize;
   int64_t mask_idx = vec_idx / batch_size;
   using VecType = kps::details::VectorType<T, VecSize>;
@@ -135,8 +140,8 @@ void MaskedFillKernel(const Context& dev_ctx,
     return;
   }
 
-  auto x_dims = x.dims();
-  auto mask_dims = mask.dims();
+  const auto& x_dims = x.dims();
+  const auto& mask_dims = mask.dims();
 
   auto expanded_size =
       common::vectorize(phi::funcs::BroadcastTwoDims(x_dims, mask_dims, -1));
@@ -169,8 +174,6 @@ void MaskedFillKernel(const Context& dev_ctx,
 
   out->Resize(expanded_dims);
 
-  auto input_dim = x_expand.dims();
-  auto mask_dim = mask_expand.dims();
   GPUMaskedFill<T>(dev_ctx, x_expand, mask_expand, value, out);
 }
 
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -2114,6 +2114,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_op : masked_select_grad
   forward : masked_select (Tensor x, Tensor mask) -> Tensor(out)
diff --git a/python/paddle/tensor/manipulation.py b/python/paddle/tensor/manipulation.py
@@ -6417,8 +6417,47 @@ def masked_fill(
     if np.isscalar(value):
         value = paddle.full([], value, x.dtype)
 
-    out = _C_ops.masked_fill(x, mask, value)
-    return out
+    if in_dynamic_or_pir_mode():
+        out = _C_ops.masked_fill(x, mask, value)
+        return out
+    else:
+        check_variable_and_dtype(
+            x,
+            'x',
+            [
+                'bool',
+                'float16',
+                'bfloat16',
+                'float32',
+                'float64',
+                'int16',
+                'int32',
+                'int64',
+                'int8',
+                'unit8',
+                'complex64',
+                'complex128',
+            ],
+            'masked_fill',
+        )
+        check_variable_and_dtype(
+            mask,
+            'mask',
+            ['bool'],
+            'masked_fill',
+        )
+        helper = LayerHelper("masked_fill", **locals())
+        out = helper.create_variable_for_type_inference(dtype=x.dtype)
+        helper.append_op(
+            type='masked_fill',
+            inputs={
+                'x': x,
+                'mask': mask,
+                'value': value,
+            },
+            outputs={'out': out},
+        )
+        return out
 
 
 @inplace_apis_in_dygraph_only
diff --git a/test/legacy_test/test_masked_fill.py b/test/legacy_test/test_masked_fill.py
@@ -256,6 +256,33 @@ def init(self):
         self.scalar_value = True
 
 
+class TestMaskedFillAPIBroadcast6(TestMaskedFillAPI):
+    def init(self):
+        self.x_shape = (1, 1)
+        self.mask_shape = (40, 40)
+        self.dtype = "float32"
+        self.scalar_value = True
+
+
+class TestMaskedFillAPIBroadcast7(TestMaskedFillAPI):
+    def init(self):
+        self.x_shape = (15,)
+        self.mask_shape = (40, 1)
+        self.dtype = "float32"
+        self.scalar_value = True
+
+
+class TestMaskedFillAPIBroadcast8(TestMaskedFillAPI):
+    def init(self):
+        self.x_shape = (3, 1, 1)
+        self.mask_shape = (
+            120,
+            40,
+        )
+        self.dtype = "float32"
+        self.scalar_value = True
+
+
 @unittest.skipIf(
     not core.is_compiled_with_cuda(), "core is not compiled with CUDA"
 )
