[Comp] Support masked_fill_decomp (#73225)

* support masked_fill_decomp

* update decomp interface conf

* fix

Author: HydrogenSulfate

paddle/fluid/pir/dialect/op_generator/decomp_interface_gen_op_list.py

@@ -78,6 +78,7 @@
     "unbind",
     "unsqueeze",
     "unstack",
+    "masked_fill",
 ]
 decomp_rule_interface_declare_gen_op_list = (
     GENERATE_IMPL_DECOMP + MANUAL_IMPL_DECOMP

paddle/fluid/primitive/decomp_rule/decomp_rule/composite.h

@@ -600,6 +600,38 @@ Tensor full_like_decomp(const Tensor& x,
   }
 }
 
+template <typename T>
+Tensor masked_fill_decomp(const Tensor& x,
+                          const Tensor& mask,
+                          const Tensor& v) {
+  if (has_dynamic_shape(x.shape()) || has_dynamic_shape(mask.shape())) {
+    // NOTE: use add operator to get broadcast shape implicitly,
+    // which is not efficient yet, should be improved in the future.
+    Tensor dummy_x =
+        backend::full_with_tensor<T>(shape64<T>(x), 0.0, x.dtype(), x.place());
+    Tensor dummy_y = backend::full_with_tensor<T>(
+        shape64<T>(mask), 0.0, x.dtype(), x.place());
+    Tensor dummy = dummy_x + dummy_y;
+    Tensor mask_expanded = backend::expand<T>(mask, shape64<T>(dummy));
+    Tensor v_expanded = backend::expand<T>(v, shape64<T>(dummy));
+    return where<T>(mask_expanded, v_expanded, x);
+
+  } else {
+    auto out_dims = phi::funcs::BroadcastTwoDims(x.dims(), mask.dims());
+    std::vector<int64_t> out_shape = common::vectorize(out_dims);
+    Tensor x_expanded = x;
+    if (x.dims() != out_dims) {
+      x_expanded = expand<T>(x_expanded, out_shape);
+    }
+    Tensor mask_expanded = mask;
+    if (mask.dims() != out_dims) {
+      mask_expanded = expand<T>(mask, out_shape);
+    }
+    Tensor v_expanded = expand<T>(v, out_shape);
+    return where<T>(mask_expanded, v_expanded, x_expanded);
+  }
+}
+
 template <typename T>
 std::tuple<Tensor, Tensor> dropout_decomp(
     const Tensor& x,