wehub-resource-sync
274 lines
10 KiB
Plaintext
274 lines
10 KiB
Plaintext
// 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 <climits>
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/common/scalar.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/tensor_utils.h"
|
|
#include "paddle/phi/kernels/index_fill_kernel.h"
|
|
|
|
namespace phi {
|
|
|
|
// GPU kernel for index_fill forward pass.
|
|
//
|
|
// Core idea — "Three-Segment Decomposition":
|
|
// For a tensor with shape [d0, d1, ..., d_{dim}, ..., d_{n-1}], we split
|
|
// the dimensions into three groups around the target `dim`:
|
|
//
|
|
// outer_size = d0 * d1 * ... * d_{dim-1} (all dims before `dim`)
|
|
// dim_size = d_{dim} (the target dim itself)
|
|
// inner_size = d_{dim+1} * ... * d_{n-1} (all dims after `dim`)
|
|
//
|
|
// This lets us treat ANY N-dimensional tensor as a 3D logical block
|
|
// [outer_size, dim_size, inner_size], and locate any element with:
|
|
//
|
|
// offset = outer_idx * (dim_size * inner_size) + dim_idx * inner_size +
|
|
// inner_idx
|
|
//
|
|
// Total threads = outer_size * index_size * inner_size
|
|
// (one thread per element that needs to be filled)
|
|
//
|
|
// IndexT: int32_t when numel <= INT32_MAX (faster mod/div on GPU),
|
|
// int64_t otherwise.
|
|
// IndT: matches index tensor dtype (int32_t or int64_t), avoids a
|
|
// CastToInt64 kernel launch when the user supplies int32 indices.
|
|
template <typename T, typename IndexT, typename IndT>
|
|
__global__ void IndexFillCudaKernel(const T* x,
|
|
const IndT* index,
|
|
const IndexT index_size,
|
|
const int dim,
|
|
const IndexT outer_size,
|
|
const int64_t dim_size,
|
|
const IndexT inner_size,
|
|
const T fill_value,
|
|
T* out) {
|
|
IndexT idx =
|
|
static_cast<IndexT>(threadIdx.x) +
|
|
static_cast<IndexT>(blockIdx.x) * static_cast<IndexT>(blockDim.x);
|
|
IndexT total = index_size * outer_size * inner_size;
|
|
if (idx >= total) return;
|
|
|
|
// Decompose the flat thread index into the three logical coordinates.
|
|
// The iteration order is: outer (slowest) → index → inner (fastest).
|
|
IndexT inner_idx = idx % inner_size;
|
|
IndexT temp = idx / inner_size;
|
|
IndexT index_idx = temp % index_size;
|
|
IndexT outer_idx = temp / index_size;
|
|
|
|
// Look up the actual position along the target dimension from the index
|
|
// tensor. Widen to int64 for the bounds check against dim_size.
|
|
int64_t dim_idx = static_cast<int64_t>(index[index_idx]);
|
|
if (dim_idx < 0) dim_idx += dim_size; // support negative indexing
|
|
|
|
if (dim_idx < 0 || dim_idx >= dim_size) return; // out-of-bounds guard
|
|
|
|
// Convert the 3D logical coordinate back to a flat memory offset.
|
|
IndexT offset = outer_idx * static_cast<IndexT>(dim_size) * inner_size +
|
|
static_cast<IndexT>(dim_idx) * inner_size + inner_idx;
|
|
|
|
out[offset] = fill_value;
|
|
}
|
|
|
|
// Host-side launch function: computes the three-segment sizes and launches
|
|
// the CUDA kernel with the appropriate index types (IndexT, IndT).
|
|
template <typename T, typename Context, typename IndexT, typename IndT>
|
|
void LaunchIndexFillCudaKernelImpl(const Context& dev_ctx,
|
|
const T* x_data,
|
|
const IndT* index_data,
|
|
IndexT index_size,
|
|
int dim,
|
|
IndexT outer_size,
|
|
int64_t dim_size,
|
|
IndexT inner_size,
|
|
T fill_value,
|
|
T* out_data) {
|
|
IndexT numel = outer_size * index_size * inner_size;
|
|
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, numel);
|
|
IndexFillCudaKernel<T, IndexT, IndT>
|
|
<<<config.block_per_grid, config.thread_per_block, 0, dev_ctx.stream()>>>(
|
|
x_data,
|
|
index_data,
|
|
index_size,
|
|
dim,
|
|
outer_size,
|
|
dim_size,
|
|
inner_size,
|
|
fill_value,
|
|
out_data);
|
|
}
|
|
|
|
template <typename T, typename Context, typename IndT>
|
|
void LaunchIndexFillCudaKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
int dim,
|
|
const DenseTensor& index,
|
|
const Scalar& value,
|
|
DenseTensor* out) {
|
|
auto* x_data = x.data<T>();
|
|
T fill_value = value.to<T>();
|
|
|
|
// "Copy-then-modify" pattern: first copy x entirely into out, then
|
|
// overwrite only the positions specified by the index.
|
|
// Skip the copy if out already shares memory with x (inplace mode).
|
|
bool is_initialized = out->initialized();
|
|
T* out_data = dev_ctx.template Alloc<T>(out);
|
|
if (!is_initialized || (x.data<T>() != out->data<T>())) {
|
|
Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
|
|
}
|
|
|
|
const IndT* index_data = index.data<IndT>();
|
|
int64_t index_size = index.numel();
|
|
|
|
if (index_size == 0) {
|
|
return;
|
|
}
|
|
|
|
// --- Three-segment decomposition ---
|
|
auto x_dims = x.dims();
|
|
const int rank = x_dims.size();
|
|
|
|
if (dim < 0) {
|
|
dim += rank;
|
|
}
|
|
|
|
int64_t outer_size = 1; // product of dims before `dim`
|
|
int64_t inner_size = 1; // product of dims after `dim`
|
|
int64_t dim_size = x_dims[dim];
|
|
|
|
for (int i = 0; i < dim; ++i) {
|
|
outer_size *= x_dims[i];
|
|
}
|
|
for (int i = dim + 1; i < rank; ++i) {
|
|
inner_size *= x_dims[i];
|
|
}
|
|
|
|
// Select int32 index arithmetic when the total work fits in 32 bits —
|
|
// GPU mod/div on 32-bit integers is significantly faster than on 64-bit.
|
|
const int64_t numel = x.numel();
|
|
constexpr int64_t kInt32Max = static_cast<int64_t>(INT32_MAX);
|
|
|
|
if (numel <= kInt32Max) {
|
|
LaunchIndexFillCudaKernelImpl<T, Context, int32_t, IndT>(
|
|
dev_ctx,
|
|
x_data,
|
|
index_data,
|
|
static_cast<int32_t>(index_size),
|
|
dim,
|
|
static_cast<int32_t>(outer_size),
|
|
dim_size,
|
|
static_cast<int32_t>(inner_size),
|
|
fill_value,
|
|
out_data);
|
|
} else {
|
|
LaunchIndexFillCudaKernelImpl<T, Context, int64_t, IndT>(dev_ctx,
|
|
x_data,
|
|
index_data,
|
|
index_size,
|
|
dim,
|
|
outer_size,
|
|
dim_size,
|
|
inner_size,
|
|
fill_value,
|
|
out_data);
|
|
}
|
|
}
|
|
|
|
// Top-level kernel entry: validates inputs and dispatches to the launcher.
|
|
template <typename T, typename Context>
|
|
void IndexFillKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const DenseTensor& index,
|
|
int dim,
|
|
const Scalar& value,
|
|
DenseTensor* out) {
|
|
// Early return for zero-element output tensor.
|
|
if (out && out->numel() == 0) {
|
|
dev_ctx.template Alloc<T>(out);
|
|
return;
|
|
}
|
|
|
|
auto x_dims = x.dims();
|
|
const int rank = x_dims.size();
|
|
|
|
// Normalize negative dim and validate range.
|
|
int real_dim = dim;
|
|
if (real_dim < 0) {
|
|
real_dim += rank;
|
|
}
|
|
|
|
PADDLE_ENFORCE_GE(real_dim,
|
|
0,
|
|
common::errors::InvalidArgument(
|
|
"The dim must be >= -%d and < %d, but received %d.",
|
|
rank,
|
|
rank,
|
|
dim));
|
|
PADDLE_ENFORCE_LT(real_dim,
|
|
rank,
|
|
common::errors::InvalidArgument(
|
|
"The dim must be >= -%d and < %d, but received %d.",
|
|
rank,
|
|
rank,
|
|
dim));
|
|
|
|
// index_fill only supports 1-D index tensors (a list of positions along 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()));
|
|
|
|
// Empty index means nothing to fill; just copy x to out.
|
|
if (index.numel() == 0) {
|
|
Copy(dev_ctx, x, dev_ctx.GetPlace(), false, out);
|
|
return;
|
|
}
|
|
|
|
// Dispatch directly on index dtype — no cast kernel needed.
|
|
if (index.dtype() == DataType::INT32) {
|
|
LaunchIndexFillCudaKernel<T, Context, int32_t>(
|
|
dev_ctx, x, real_dim, index, value, out);
|
|
} else if (index.dtype() == DataType::INT64) {
|
|
LaunchIndexFillCudaKernel<T, Context, int64_t>(
|
|
dev_ctx, x, real_dim, index, value, out);
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The dtype of index must be int32 or int64, but received %s.",
|
|
DataTypeToString(index.dtype())));
|
|
}
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(index_fill,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::IndexFillKernel,
|
|
float,
|
|
double,
|
|
int,
|
|
int64_t,
|
|
bool,
|
|
int16_t,
|
|
uint8_t,
|
|
int8_t,
|
|
phi::float16,
|
|
phi::bfloat16,
|
|
phi::complex64,
|
|
phi::complex128) {}
|