Add index_fill / index_fill_ operator with CPU/GPU kernels (PaddlePaddle#78041)

* index_fill writing

* solve static gra

* delete annotations

* testfile change

Author: xpislp

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.cc

@@ -2683,6 +2683,31 @@ bool IndexAdd_OpInferSymbolicShape(
   return IndexAddOpInferSymbolicShape(op, infer_context);
 }
 
+bool IndexFillOpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  const auto &x_shape_or_data =
+      infer_context->GetShapeOrDataForValue(op->operand_source(0));
+  std::vector<symbol::DimExpr> x_shape = x_shape_or_data.shape();
+
+  PADDLE_ENFORCE_LT(
+      x_shape.size(),
+      7,
+      common::errors::InvalidArgument(
+          "The rank of input should be less than 7, but received %d.",
+          x_shape.size()));
+
+  infer_context->SetShapeOrDataForValue(
+      op->result(0),
+      symbol::ShapeOrDataDimExprs{symbol::TensorShapeOrDataDimExprs(x_shape)});
+
+  return true;
+}
+
+bool IndexFill_OpInferSymbolicShape(
+    pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
+  return IndexFillOpInferSymbolicShape(op, infer_context);
+}
+
 bool IndexPutOpInferSymbolicShape(
     pir::Operation *op, pir::InferSymbolicShapeContext *infer_context) {
   const auto &x_shape_or_data =

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/binary_infer_sym.h

@@ -63,6 +63,8 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Histogram)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Isclose)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexAdd_)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexFill)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexFill_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexPut_)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(IndexSelect)

paddle/phi/infermeta/multiary.cc

@@ -3813,6 +3813,50 @@ void LegacyInterpolateInferMeta(
                        config);
 }
 
+void IndexFillInferMeta(const MetaTensor& x,
+                        const MetaTensor& index,
+                        int dim,
+                        const Scalar& value,
+                        MetaTensor* out) {
+  auto in_dims = x.dims();
+  auto index_dims = index.dims();
+  int rank = in_dims.size();
+
+  PADDLE_ENFORCE_LT(
+      rank,
+      7,
+      common::errors::InvalidArgument(
+          "The rank of Input(X) should be less than 7, but received %d.",
+          rank));
+
+  if (dim < 0) {
+    dim += rank;
+  }
+
+  PADDLE_ENFORCE_GE(dim,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The dim must be >= -%d and < %d, but received %d.",
+                        rank,
+                        rank,
+                        dim));
+  PADDLE_ENFORCE_LT(dim,
+                    rank,
+                    common::errors::InvalidArgument(
+                        "The dim must be >= -%d and < %d, but received %d.",
+                        rank,
+                        rank,
+                        dim));
+
+  PADDLE_ENFORCE_EQ(index_dims.size(),
+                    1,
+                    common::errors::InvalidArgument(
+                        "The index tensor must be 1-D, but received %d-D.",
+                        index_dims.size()));
+
+  out->share_meta(x);
+}
+
 void IndexPutInferMeta(const MetaTensor& x,
                        const std::vector<const MetaTensor*>& indices,
                        const MetaTensor& value,

paddle/phi/infermeta/multiary.h

@@ -739,6 +739,12 @@ PADDLE_API void LegacyInterpolateInferMeta(
     MetaTensor* output,
     MetaConfig config = MetaConfig());
 
+PADDLE_API void IndexFillInferMeta(const MetaTensor& x,
+                                   const MetaTensor& index,
+                                   int dim,
+                                   const Scalar& value,
+                                   MetaTensor* out);
+
 PADDLE_API void IndexPutInferMeta(const MetaTensor& x,
                                   const std::vector<const MetaTensor*>& indices,
                                   const MetaTensor& value,

paddle/phi/kernels/cpu/index_fill_grad_kernel.cc

@@ -0,0 +1,170 @@
+// Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/index_fill_grad_kernel.h"
+#include "paddle/phi/backends/cpu/cpu_context.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/tensor_utils.h"
+
+namespace phi {
+
+// CPU implementation of the index_fill backward kernel.
+// Same logic as the GPU version:
+//   For each position selected by the index, set x_grad to 0.
+//
+// The flat iteration index is decomposed into (outer, index, inner) using
+// the same three-segment scheme, with OMP parallelization on the outermost
+// loop.
+template <typename T>
+void index_fill_grad_kernel(const int64_t N,
+                            const int64_t* index_data,
+                            const int64_t index_size,
+                            const int64_t dim_size,
+                            const int64_t outer_size,
+                            const int64_t inner_size,
+                            T* x_grad) {
+#ifdef PADDLE_WITH_MKLML
+#pragma omp parallel for
+#endif
+  for (int64_t idx = 0; idx < N; ++idx) {
+    // Decompose flat index → (outer_idx, index_idx, inner_idx)
+    int64_t inner_idx = idx % inner_size;
+    int64_t temp = idx / inner_size;
+    int64_t index_idx = temp % index_size;
+    int64_t outer_idx = temp / index_size;
+
+    int64_t dim_idx = index_data[index_idx];
+    if (dim_idx < 0) {
+      dim_idx += dim_size;
+    }
+
+    int64_t offset =
+        outer_idx * dim_size * inner_size + dim_idx * inner_size + inner_idx;
+
+    // Zero out gradient at the filled position.
+    *(x_grad + offset) = static_cast<T>(0);
+  }
+}
+
+// CPU host-side launch function for the backward kernel.
+template <typename T, typename Context>
+void LaunchIndexFillGradKernel(const Context& dev_ctx,
+                               const DenseTensor& index,
+                               const DenseTensor& out_grad,
+                               const int dim,
+                               DenseTensor* x_grad) {
+  // Step 1: x_grad = out_grad (full copy).
+  T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
+  Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+
+  auto out_grad_dims = out_grad.dims();
+  const int rank = out_grad_dims.size();
+
+  // Cast index to int64 if needed.
+  DenseTensor index_int64;
+  const DenseTensor* ptr_index = nullptr;
+
+  if (index.dtype() == phi::DataType::INT32) {
+    index_int64.Resize(index.dims());
+    int64_t* index_int64_data = dev_ctx.template Alloc<int64_t>(&index_int64);
+    const int32_t* index_int32_data = index.data<int32_t>();
+
+    int64_t index_numel = index.numel();
+    for (int64_t i = 0; i < index_numel; ++i) {
+      index_int64_data[i] = static_cast<int64_t>(index_int32_data[i]);
+    }
+
+    ptr_index = &index_int64;
+  } else {
+    ptr_index = &index;
+  }
+
+  const int64_t* index_data = ptr_index->data<int64_t>();
+  int64_t index_size = ptr_index->numel();
+
+  if (index_size == 0) {
+    return;
+  }
+
+  // Three-segment decomposition (same as forward).
+  int64_t outer_size = 1;
+  int64_t inner_size = 1;
+  int64_t dim_size = out_grad_dims[dim];
+
+  for (int i = 0; i < dim; ++i) {
+    outer_size *= out_grad_dims[i];
+  }
+  for (int i = dim + 1; i < rank; ++i) {
+    inner_size *= out_grad_dims[i];
+  }
+
+  // Step 2: zero out the positions that were filled in forward pass.
+  int64_t numel = outer_size * index_size * inner_size;
+
+  index_fill_grad_kernel<T>(numel,
+                            index_data,
+                            index_size,
+                            dim_size,
+                            outer_size,
+                            inner_size,
+                            x_grad_data);
+}
+
+// Top-level CPU backward kernel entry.
+template <typename T, typename Context>
+void IndexFillGradKernel(const Context& dev_ctx,
+                         const DenseTensor& index,
+                         const DenseTensor& out_grad,
+                         int dim,
+                         DenseTensor* x_grad) {
+  if (out_grad.numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
+
+  dev_ctx.template Alloc<T>(x_grad);
+
+  auto out_grad_dims = out_grad.dims();
+  const int rank = out_grad_dims.size();
+
+  if (dim < 0) {
+    dim += rank;
+  }
+
+  if (index.numel() == 0) {
+    Copy(dev_ctx, out_grad, dev_ctx.GetPlace(), false, x_grad);
+    return;
+  }
+
+  LaunchIndexFillGradKernel<T, Context>(dev_ctx, index, out_grad, dim, x_grad);
+}
+
+}  // namespace phi
+
+PD_REGISTER_KERNEL(index_fill_grad,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::IndexFillGradKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t,
+                   bool,
+                   int16_t,
+                   uint8_t,
+                   int8_t,
+                   phi::float16,
+                   phi::bfloat16,
+                   phi::complex64,
+                   phi::complex128) {}