Delete Index64 implement in broadcast func (#57996)

AnnaTrainingG · web-flow · commit 250340858a70 · 2023-10-11T11:55:10.000+08:00
diff --git a/paddle/phi/kernels/funcs/broadcast_function.h b/paddle/phi/kernels/funcs/broadcast_function.h
@@ -511,339 +511,13 @@ void LaunchBroadcastKernel(
 #endif
 }
 
-#ifndef PADDLE_WITH_XPU_KP
-HOSTDEVICE static int64_t ConvertSrcIdxToDstIdx(
-    int64_t src_idx,
-    const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &src_strides,
-    const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &dst_strides,
-    int rank) {
-  int64_t dst_idx = 0;
-  int64_t origin_src_idx = src_idx;
-  for (int k = 0; k < rank; ++k) {
-    auto local_idx = src_idx / src_strides[k + 1];
-    src_idx -= local_idx * src_strides[k + 1];
-
-    if (dst_strides[k] != dst_strides[k + 1]) {
-      dst_idx += local_idx * dst_strides[k + 1];
-    }
-  }
-  return dst_idx;
-}
-
-template <int Index, int VecSize>
-struct ReadVecDataWithInt64Index {
-  template <typename Array1, typename Array2, typename Array3, typename ArgsT>
-  static __device__ __forceinline__ void Apply(
-      const Array1 &in,
-      ArgsT *args,
-      int64_t idx,
-      const Array2 &need_broadcast,
-      const phi::Array<int64_t, phi::DDim::kMaxRank + 1> &src_strides,
-      const Array3 &dst_strides,
-      int rank,
-      bool is_boundary) {
-    using Type = std::tuple_element_t<Index, ArgsT>;
-    if (is_boundary) {
-#pragma unroll
-      for (int i = 0; i < VecSize; ++i) {
-        std::get<Index>(args[i]) = in[Index][ConvertSrcIdxToDstIdx(
-            idx + i, src_strides, dst_strides[Index], rank)];
-      }
-    } else {
-      if (!need_broadcast[Index]) {
-        kps::ReadData<Type, VecSize, 1, ArgsT, Index, false>(
-            args, reinterpret_cast<const _ptr_ Type *>(in[Index]) + idx, 1);
-      } else {
-#pragma unroll
-        for (int i = 0; i < VecSize; ++i) {
-          std::get<Index>(args[i]) = in[Index][ConvertSrcIdxToDstIdx(
-              idx + i, src_strides, dst_strides[Index], rank)];
-        }
-      }
-    }
-  }
-};
-
-template <typename OutT, typename Functor, int VecSize, int NumIns>
-__global__ void BroadcastKernelWithInt64Index(
-    const phi::Array<const _ptr_ char *__restrict__, MaxWithOne<NumIns>::kValue>
-        &ins,
-    OutT *out,
-    phi::Array<phi::Array<int64_t, phi::DDim::kMaxRank + 1>,
-               MaxWithOne<NumIns>::kValue> ins_strides,
-    phi::Array<int64_t, phi::DDim::kMaxRank + 1> out_strides,
-    phi::Array<bool, MaxWithOne<NumIns>::kValue> need_broadcasts,
-    int rank,
-    Functor functor) {
-  int64_t numel = out_strides[0];
-  int64_t idx =
-      (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) * VecSize;
-  int64_t stride = static_cast<int64_t>(blockDim.x) * gridDim.x * VecSize;
-  int64_t limit = numel - VecSize;
-
-  using Traits = phi::funcs::FunctionTraits<Functor>;
-  using ArgsT = typename Traits::ArgsTuple;
-
-  ArgsT args[VecSize];
-  phi::AlignedVector<OutT, VecSize> out_vec;
-  for (; idx <= limit; idx += stride) {
-    Unroller<ReadVecDataWithInt64Index, VecSize, NumIns>::step(
-        ins, args, idx, need_broadcasts, out_strides, ins_strides, rank, false);
-
-#pragma unroll
-    for (int i = 0; i < VecSize; ++i) {
-      out_vec[i] = static_cast<OutT>(Apply(functor, args[i]));
-    }
-    phi::Store<OutT, VecSize>(out_vec, out + idx);
-  }
-
-  if (idx < numel) {
-    int remain = numel - idx;  // remain is always less than VecSize, therefore
-                               // `int` is enough here
-    Unroller<ReadVecDataWithInt64Index, VecSize, NumIns>::step(
-        ins, args, idx, need_broadcasts, out_strides, ins_strides, rank, true);
-    for (int i = 0; i < remain; ++i) {
-      out_vec[idx + i] = static_cast<OutT>(Apply(functor, args[i]));
-    }
-  }
-}
-
-template <typename OutT, typename Functor, int Arity, int NumOuts, int VecSize>
-struct LaunchBroadcastKernelWithInt64IndexHelper {
-  static void Run(const KPDevice &ctx,
-                  const std::vector<const DenseTensor *> &ins,
-                  std::vector<DenseTensor *> *outs,
-                  int axis,
-                  Functor functor) {
-    PADDLE_THROW(phi::errors::PermissionDenied(
-        "Unreachable code branch. This may be a bug."));
-  }
-};
-
-template <typename OutT, typename Functor, int Arity, int VecSize>
-struct LaunchBroadcastKernelWithInt64IndexHelper<OutT,
-                                                 Functor,
-                                                 Arity,
-                                                 /*NumOuts=*/1,
-                                                 VecSize> {
-  static void Run(const KPDevice &ctx,
-                  const std::vector<const DenseTensor *> &ins,
-                  std::vector<DenseTensor *> *outs,
-                  int axis,
-                  Functor functor) {
-    using Traits = phi::funcs::FunctionTraits<Functor>;
-    using ArgsT = typename Traits::ArgsTuple;
-    ArgsT arg;
-    phi::Array<const _ptr_ char *__restrict__, MaxWithOne<Arity>::kValue>
-        ins_ptrs;
-    UnrollerWithoutVecSize<InputSetter, Arity>::step(ins, arg, &ins_ptrs);
-
-    auto *out_tensor = (*outs)[0];
-    auto *out_ptr = ctx.Alloc<OutT>(out_tensor);
-
-    phi::Array<phi::Array<int64_t, phi::DDim::kMaxRank>,
-               MaxWithOne<Arity>::kValue>
-        ins_expand_dims;
-    phi::Array<int64_t, phi::DDim::kMaxRank> broadcast_out_dims;
-    int rank;
-    if (Arity == 1) {
-      rank = ins[0]->dims().size();
-      for (int i = 0; i < rank; ++i) {
-        broadcast_out_dims[i] = ins[0]->dims()[i];
-      }
-      ins_expand_dims[0] = broadcast_out_dims;
-    } else if (Arity >= 2) {
-      CalculateBroadcastDims(ins[0]->dims().Get(),
-                             ins[1]->dims().Get(),
-                             ins[0]->dims().size(),
-                             ins[1]->dims().size(),
-                             axis,
-                             ins_expand_dims[0].GetMutable(),
-                             ins_expand_dims[1].GetMutable(),
-                             broadcast_out_dims.GetMutable(),
-                             &rank);
-      for (int i = 2; i < Arity; ++i) {
-        auto tmp_dims = broadcast_out_dims;
-        phi::Array<int64_t, phi::DDim::kMaxRank> tmp_expand_dims;
-        int tmp_rank;
-        PADDLE_ENFORCE_GE(rank,
-                          ins[i]->dims().size(),
-                          phi::errors::InvalidArgument(
-                              "Unsupported reverse broadcast when the input "
-                              "tensor number is larger than 2."));
-        CalculateBroadcastDims(tmp_dims.Get(),
-                               ins[i]->dims().Get(),
-                               rank,
-                               ins[i]->dims().size(),
-                               axis,
-                               tmp_expand_dims.GetMutable(),
-                               ins_expand_dims[i].GetMutable(),
-                               broadcast_out_dims.GetMutable(),
-                               &tmp_rank);
-        PADDLE_ENFORCE_EQ(rank,
-                          tmp_rank,
-                          phi::errors::InvalidArgument(
-                              "Wrong broadcast algorithm. This may be a bug."));
-      }
-    }
-
-    phi::Array<phi::Array<int64_t, phi::DDim::kMaxRank + 1>,
-               MaxWithOne<Arity>::kValue>
-        ins_strides;
-    phi::Array<bool, MaxWithOne<Arity>::kValue> need_broadcasts;
-    phi::Array<int64_t, phi::DDim::kMaxRank + 1> out_strides;
-    const auto &out_dims = out_tensor->dims();
-    if (rank <= out_dims.size()) {
-      out_strides = ShapeToStride(out_dims.Get(), rank);
-    } else {
-      out_strides = ShapeToStride(broadcast_out_dims.Get(), rank);
-    }
-
-    for (int i = 0; i < Arity; ++i) {
-      ins_strides[i] = ShapeToStride(ins_expand_dims[i].Get(), rank);
-      need_broadcasts[i] =
-          !IsSameShape(out_strides.Get(), ins_strides[i].Get(), rank + 1);
-    }
-
-    int64_t numel = out_strides[0];
-    auto gpu_config =
-        phi::backends::gpu::GetGpuLaunchConfig1D(ctx, numel, VecSize);
-
-    BroadcastKernelWithInt64Index<OutT, Functor, VecSize, Arity>
-        <<<gpu_config.block_per_grid,
-           gpu_config.thread_per_block,
-           0,
-           ctx.stream()>>>(ins_ptrs,
-                           out_ptr,
-                           ins_strides,
-                           out_strides,
-                           need_broadcasts,
-                           rank,
-                           functor);
-  }
-
- private:
-  static void CalculateBroadcastDims(const int64_t *x_dims,
-                                     const int64_t *y_dims,
-                                     int nx,
-                                     int ny,
-                                     int axis,
-                                     int64_t *x_out_dims,
-                                     int64_t *y_out_dims,
-                                     int64_t *broadcast_out_dims,
-                                     int *length) {
-    PADDLE_ENFORCE_GE(
-        axis, 0, phi::errors::InvalidArgument("Invalid axis value: %d", axis));
-    if (nx == ny) {
-      *length = nx;
-      for (int i = 0; i < nx; ++i) {
-        if (x_dims[i] != y_dims[i]) {
-          PADDLE_ENFORCE_EQ(
-              x_dims[i] == 1 || y_dims[i] == 1,
-              true,
-              phi::errors::InvalidArgument("Cannot broadcast input shape where "
-                                           "x_dims[%d] = %d, y_dims[%d] = %d.",
-                                           i,
-                                           x_dims[i],
-                                           i,
-                                           y_dims[i]));
-        }
-        broadcast_out_dims[i] = std::max(x_dims[i], y_dims[i]);
-        x_out_dims[i] = x_dims[i];
-        y_out_dims[i] = y_dims[i];
-      }
-    } else if (nx > ny) {
-      *length = nx;
-      for (int i = nx - axis; i < ny; ++i) {
-        PADDLE_ENFORCE_EQ(
-            y_dims[i],
-            1,
-            phi::errors::InvalidArgument(
-                "The trailing Y.shape[%d] should be 1 but got %d.",
-                i,
-                y_dims[i]));
-      }
-
-      for (int i = 0; i < nx; ++i) {
-        if (i >= axis && i - axis < ny) {
-          if (x_dims[i] != y_dims[i - axis]) {
-            PADDLE_ENFORCE_EQ(x_dims[i] == 1 || y_dims[i - axis] == 1,
-                              true,
-                              phi::errors::InvalidArgument(
-                                  "Cannot broadcast input shape where "
-                                  "x_dims[%d] = %d, y_dims[%d] = %d.",
-                                  i,
-                                  x_dims[i],
-                                  i - axis,
-                                  y_dims[i - axis]));
-          }
-          broadcast_out_dims[i] = std::max(x_dims[i], y_dims[i - axis]);
-          x_out_dims[i] = x_dims[i];
-          y_out_dims[i] = y_dims[i - axis];
-        } else {
-          broadcast_out_dims[i] = x_dims[i];
-          x_out_dims[i] = x_dims[i];
-          y_out_dims[i] = 1;
-        }
-      }
-    } else {
-      CalculateBroadcastDims(y_dims,
-                             x_dims,
-                             ny,
-                             nx,
-                             axis,
-                             y_out_dims,
-                             x_out_dims,
-                             broadcast_out_dims,
-                             length);
-    }
-  }
-
-  static bool IsSameShape(const int64_t *x, const int64_t *y, int rank) {
-    for (int i = 0; i < rank; ++i) {
-      if (x[i] != y[i]) return false;
-    }
-    return true;
-  }
-
-  static phi::Array<int64_t, phi::DDim::kMaxRank + 1> ShapeToStride(
-      const int64_t *arr, int rank) {
-    phi::Array<int64_t, phi::DDim::kMaxRank + 1> strides;
-    strides[rank] = 1;
-    for (int i = rank - 1; i >= 0; --i) {
-      strides[i] = strides[i + 1] * arr[i];
-    }
-    return strides;
-  }
-};
-#endif
-
 template <typename OutT, typename Functor, int Arity, int NumOuts = 1>
 typename std::enable_if<!NeedVectorized<OutT>::value, void>::type
 BroadcastKernelForDifferentVecSize(const KPDevice &ctx,
                                    const std::vector<const DenseTensor *> &ins,
                                    std::vector<DenseTensor *> *outs,
                                    int axis,
                                    Functor func) {
-#ifndef PADDLE_WITH_XPU_KP
-  constexpr bool kEnabledInt64IndexKernel = (NumOuts == 1 && Arity <= 3);
-  bool use_int64_index_kernel =
-      kEnabledInt64IndexKernel &&
-      (*outs)[0]->numel() >= std::numeric_limits<int32_t>::max();
-  if (use_int64_index_kernel) {
-    LaunchBroadcastKernelWithInt64IndexHelper<OutT,
-                                              Functor,
-                                              Arity,
-                                              NumOuts,
-                                              VecSizeS>::Run(ctx,
-                                                             ins,
-                                                             outs,
-                                                             axis,
-                                                             func);
-    return;
-  }
-#endif
-
   auto classifier =
       BroadcastTypeClassifier<OutT, Functor, Arity, NumOuts>(ins, outs, axis);
   LaunchBroadcastKernel<OutT, Functor, Arity, NumOuts, VecSizeS>(
