chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,579 @@
|
||||
/* Copyright (c) 2022 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 <thrust/execution_policy.h>
|
||||
|
||||
#include "paddle/phi/kernels/sparse/mask_kernel.h"
|
||||
#include "paddle/phi/kernels/sparse/sparse_utils_kernel.h"
|
||||
|
||||
#include "paddle/common/ddim.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_info.h"
|
||||
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
#include "paddle/phi/core/visit_type.h"
|
||||
#include "paddle/phi/kernels/empty_kernel.h"
|
||||
#include "paddle/phi/kernels/funcs/aligned_vector.h"
|
||||
#include "paddle/phi/kernels/funcs/math_function.h"
|
||||
#include "paddle/phi/kernels/funcs/sparse/flatten_indices.cu.h"
|
||||
#include "paddle/phi/kernels/funcs/sparse/utils.cu.h"
|
||||
|
||||
namespace phi {
|
||||
namespace sparse {
|
||||
|
||||
template <typename T, typename IntT>
|
||||
__global__ void MaskKernel(const T* x_ptr,
|
||||
const IntT* indices_ptr,
|
||||
const int64_t* sparse_offsets,
|
||||
const int64_t non_zero_num,
|
||||
const int cols,
|
||||
const int sparse_dim,
|
||||
T* out_values_ptr) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, non_zero_num * cols, int64_t) {
|
||||
int64_t out_i = i / cols;
|
||||
int64_t col_i = i - out_i * cols;
|
||||
int64_t index = 0;
|
||||
for (int j = 0; j < sparse_dim; j++) {
|
||||
index += indices_ptr[j * non_zero_num + out_i] * sparse_offsets[j];
|
||||
}
|
||||
out_values_ptr[out_i * cols + col_i] = x_ptr[index * cols + col_i];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT>
|
||||
void MaskCooGPUKernel(const GPUContext& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const SparseCooTensor& mask,
|
||||
SparseCooTensor* out) {
|
||||
const DDim& dims = x.dims();
|
||||
PADDLE_ENFORCE_EQ(x.dims(),
|
||||
mask.dims(),
|
||||
common::errors::InvalidArgument(
|
||||
"the input x and mask must have the shape"));
|
||||
const DenseTensor& indices = mask.indices();
|
||||
const DenseTensor& values = mask.values();
|
||||
DenseTensor out_indices = EmptyLike<IntT>(dev_ctx, indices);
|
||||
DenseTensor out_values = EmptyLike<T>(dev_ctx, values);
|
||||
if (mask.nnz() <= 0) {
|
||||
out->SetMember(out_indices, out_values, dims, true);
|
||||
return;
|
||||
}
|
||||
|
||||
const int sparse_dim = mask.sparse_dim();
|
||||
DenseTensor sparse_offsets = Empty<GPUContext>(
|
||||
dev_ctx,
|
||||
DenseTensorMeta(DataType::INT64, {sparse_dim}, DataLayout::NCHW));
|
||||
std::vector<int64_t> h_sparse_offsets(sparse_dim);
|
||||
funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
|
||||
|
||||
backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
|
||||
&h_sparse_offsets[0],
|
||||
sizeof(int64_t) * sparse_dim,
|
||||
gpuMemcpyHostToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
phi::Copy(dev_ctx, indices, dev_ctx.GetPlace(), false, &out_indices);
|
||||
|
||||
const IntT* indices_ptr = indices.data<IntT>();
|
||||
T* out_values_ptr = out_values.data<T>();
|
||||
const T* x_ptr = x.data<T>();
|
||||
const int64_t non_zero_num = mask.nnz();
|
||||
auto dims_2d = flatten_to_2d(dims, sparse_dim);
|
||||
const int cols = dims_2d[1];
|
||||
|
||||
auto config =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
|
||||
MaskKernel<T, IntT>
|
||||
<<<config.block_per_grid, config.thread_per_block, 0, dev_ctx.stream()>>>(
|
||||
x_ptr,
|
||||
indices_ptr,
|
||||
sparse_offsets.data<int64_t>(),
|
||||
non_zero_num,
|
||||
cols,
|
||||
sparse_dim,
|
||||
out_values_ptr);
|
||||
|
||||
out->SetMember(out_indices, out_values, dims, true);
|
||||
}
|
||||
|
||||
template <typename IntT>
|
||||
__global__ void ConvertCsrCrowsToCooRows(const IntT* crows_ptr,
|
||||
const IntT* crows_offsets,
|
||||
IntT* rows_ptr,
|
||||
IntT* batch_ptr,
|
||||
const int rows) {
|
||||
const int b = blockIdx.y;
|
||||
const int64_t offset = crows_offsets ? crows_offsets[b] : 0;
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
for (int i = tid; i < rows; i += gridDim.x * blockDim.x) {
|
||||
for (int j = crows_ptr[b * (rows + 1) + i];
|
||||
j < crows_ptr[b * (rows + 1) + i + 1];
|
||||
j++) {
|
||||
rows_ptr[offset + j] = i;
|
||||
if (batch_ptr) {
|
||||
batch_ptr[offset + j] = b;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IntT>
|
||||
__global__ void GetBatchSizes(const IntT* crows,
|
||||
const int rows,
|
||||
const int batches,
|
||||
IntT* batch_sizes) {
|
||||
const int tid = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
if (tid < batches) {
|
||||
batch_sizes[tid] = crows[tid * (rows + 1) + rows];
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT>
|
||||
void MaskCsr2DGPUKernel(const GPUContext& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const SparseCsrTensor& mask,
|
||||
SparseCsrTensor* out) {
|
||||
const DenseTensor& mask_cols = mask.cols();
|
||||
const DenseTensor& mask_crows = mask.crows();
|
||||
int64_t num_non_zeros = mask.nnz();
|
||||
|
||||
DenseTensor out_cols = EmptyLike<IntT>(dev_ctx, mask_cols);
|
||||
DenseTensor out_crows = EmptyLike<IntT>(dev_ctx, mask_crows);
|
||||
DenseTensor out_values = Empty<T>(dev_ctx, {num_non_zeros});
|
||||
|
||||
phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols);
|
||||
phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows);
|
||||
|
||||
const DDim& dims = x.dims();
|
||||
const int64_t non_zero_num = mask.nnz();
|
||||
int64_t sparse_dim = 2;
|
||||
DenseTensor sparse_offsets = Empty<IntT>(dev_ctx, {sparse_dim});
|
||||
std::vector<int64_t> h_sparse_offsets(sparse_dim);
|
||||
funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
|
||||
|
||||
backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
|
||||
&h_sparse_offsets[0],
|
||||
sizeof(int64_t) * sparse_dim,
|
||||
gpuMemcpyHostToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
const auto& csr_crows = mask.crows();
|
||||
const auto& csr_cols = mask.cols();
|
||||
const IntT* csr_crows_data = csr_crows.data<IntT>();
|
||||
const IntT* csr_cols_data = csr_cols.data<IntT>();
|
||||
|
||||
const int batches = 1;
|
||||
const int rows = dims[0];
|
||||
auto dims_2d = flatten_to_2d(dims, sparse_dim);
|
||||
const int cols = dims_2d[1];
|
||||
|
||||
DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
|
||||
IntT* coo_indices = indices.data<IntT>();
|
||||
IntT* batch_ptr = nullptr;
|
||||
IntT* coo_rows_data = coo_indices;
|
||||
IntT* coo_cols_data = coo_rows_data + non_zero_num;
|
||||
IntT* offsets_ptr = nullptr;
|
||||
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
|
||||
config.block_per_grid.y = batches;
|
||||
ConvertCsrCrowsToCooRows<IntT>
|
||||
<<<config.block_per_grid, config.thread_per_block.x>>>(
|
||||
csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows);
|
||||
backends::gpu::GpuMemcpyAsync(coo_cols_data,
|
||||
csr_cols_data,
|
||||
sizeof(IntT) * non_zero_num,
|
||||
gpuMemcpyDeviceToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
const T* x_ptr = x.data<T>();
|
||||
const IntT* indices_ptr = coo_indices;
|
||||
T* out_values_ptr = out_values.data<T>();
|
||||
|
||||
auto config_mask =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
|
||||
MaskKernel<T, IntT><<<config_mask.block_per_grid,
|
||||
config_mask.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(x_ptr,
|
||||
indices_ptr,
|
||||
sparse_offsets.data<int64_t>(),
|
||||
non_zero_num,
|
||||
cols,
|
||||
sparse_dim,
|
||||
out_values_ptr);
|
||||
|
||||
out->SetMember(out_crows, out_cols, out_values, x.dims());
|
||||
}
|
||||
|
||||
template <typename T, typename IntT>
|
||||
void MaskCsr3DGPUKernel(const GPUContext& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const SparseCsrTensor& mask,
|
||||
SparseCsrTensor* out) {
|
||||
const DenseTensor& mask_cols = mask.cols();
|
||||
const DenseTensor& mask_crows = mask.crows();
|
||||
int64_t num_non_zeros = mask.nnz();
|
||||
|
||||
DenseTensor out_cols = EmptyLike<IntT>(dev_ctx, mask_cols);
|
||||
DenseTensor out_crows = EmptyLike<IntT>(dev_ctx, mask_crows);
|
||||
DenseTensor out_values = Empty<T>(dev_ctx, {num_non_zeros});
|
||||
|
||||
phi::Copy(dev_ctx, mask_cols, dev_ctx.GetPlace(), false, &out_cols);
|
||||
phi::Copy(dev_ctx, mask_crows, dev_ctx.GetPlace(), false, &out_crows);
|
||||
|
||||
const DDim& dims = x.dims();
|
||||
const int64_t non_zero_num = mask.nnz();
|
||||
int64_t sparse_dim = 3;
|
||||
DenseTensor sparse_offsets = Empty<IntT>(dev_ctx, {sparse_dim});
|
||||
std::vector<int64_t> h_sparse_offsets(sparse_dim);
|
||||
funcs::sparse::CalcOffsetsPerDim(dims, sparse_dim, h_sparse_offsets.data());
|
||||
|
||||
backends::gpu::GpuMemcpyAsync(sparse_offsets.data<int64_t>(),
|
||||
&h_sparse_offsets[0],
|
||||
sizeof(int64_t) * sparse_dim,
|
||||
gpuMemcpyHostToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
const auto& csr_crows = mask.crows();
|
||||
const auto& csr_cols = mask.cols();
|
||||
const IntT* csr_crows_data = csr_crows.data<IntT>();
|
||||
const IntT* csr_cols_data = csr_cols.data<IntT>();
|
||||
|
||||
const int batches = dims[0];
|
||||
const int rows = dims[1];
|
||||
auto dims_2d = flatten_to_2d(dims, sparse_dim);
|
||||
const int cols = dims_2d[1];
|
||||
|
||||
DenseTensor indices = Empty<IntT>(dev_ctx, {sparse_dim, non_zero_num});
|
||||
DenseTensor offsets = Empty<IntT>(dev_ctx, {batches});
|
||||
IntT* coo_indices = indices.data<IntT>();
|
||||
IntT* batch_ptr = coo_indices;
|
||||
IntT* coo_rows_data = batch_ptr + non_zero_num;
|
||||
IntT* coo_cols_data = coo_rows_data + non_zero_num;
|
||||
IntT* offsets_ptr = offsets.data<IntT>();
|
||||
|
||||
auto config_batch = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, batches, 1);
|
||||
GetBatchSizes<IntT>
|
||||
<<<config_batch.block_per_grid.x, config_batch.thread_per_block.x>>>(
|
||||
csr_crows_data, rows, batches, offsets_ptr);
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
thrust::exclusive_scan(thrust::hip::par.on(dev_ctx.stream()),
|
||||
#else
|
||||
thrust::exclusive_scan(thrust::cuda::par.on(dev_ctx.stream()),
|
||||
#endif
|
||||
offsets_ptr,
|
||||
offsets_ptr + batches,
|
||||
offsets_ptr);
|
||||
|
||||
auto config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, rows, 1);
|
||||
config.block_per_grid.y = batches;
|
||||
ConvertCsrCrowsToCooRows<IntT>
|
||||
<<<config.block_per_grid, config.thread_per_block.x>>>(
|
||||
csr_crows_data, offsets_ptr, coo_rows_data, batch_ptr, rows);
|
||||
backends::gpu::GpuMemcpyAsync(coo_cols_data,
|
||||
csr_cols_data,
|
||||
sizeof(IntT) * non_zero_num,
|
||||
gpuMemcpyDeviceToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
const T* x_ptr = x.data<T>();
|
||||
const IntT* indices_ptr = coo_indices;
|
||||
T* out_values_ptr = out_values.data<T>();
|
||||
|
||||
auto config_mask =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, non_zero_num * cols, 1);
|
||||
MaskKernel<T, IntT><<<config_mask.block_per_grid,
|
||||
config_mask.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(x_ptr,
|
||||
indices_ptr,
|
||||
sparse_offsets.data<int64_t>(),
|
||||
non_zero_num,
|
||||
cols,
|
||||
sparse_dim,
|
||||
out_values_ptr);
|
||||
|
||||
out->SetMember(out_crows, out_cols, out_values, x.dims());
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Filter the DenseTensor x by the
|
||||
* mask.indices() and output a SparseCooTensor
|
||||
* x and mask must have the same shape.
|
||||
**/
|
||||
template <typename T, typename Context>
|
||||
void MaskAsCooKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const SparseCooTensor& mask,
|
||||
SparseCooTensor* out) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
mask.indices().dtype(), "MaskCooGPUKernel", ([&] {
|
||||
MaskCooGPUKernel<T, data_t>(dev_ctx, x, mask, out);
|
||||
}));
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Filter the DenseTensor x by the
|
||||
* mask.crows(), mask.cols() and output a SparseCsrTensor
|
||||
* x and mask must have the same shape.
|
||||
**/
|
||||
template <typename T, typename Context>
|
||||
void MaskAsCsrKernel(const Context& dev_ctx,
|
||||
const DenseTensor& x,
|
||||
const SparseCsrTensor& mask,
|
||||
SparseCsrTensor* out) {
|
||||
const DDim& x_dims = x.dims();
|
||||
if (x_dims.size() == 2) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
mask.crows().dtype(), "MaskCsr2DGPUKernel", ([&] {
|
||||
MaskCsr2DGPUKernel<T, data_t>(dev_ctx, x, mask, out);
|
||||
}));
|
||||
} else if (x_dims.size() == 3) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
mask.crows().dtype(), "MaskCsr3DGPUKernel", ([&] {
|
||||
MaskCsr3DGPUKernel<T, data_t>(dev_ctx, x, mask, out);
|
||||
}));
|
||||
} else {
|
||||
// throw exception
|
||||
common::errors::InvalidArgument(
|
||||
"mask_as for Sparse CSR Tensor only support 2-D or 3-D, but got "
|
||||
"%d-D.",
|
||||
x_dims.size());
|
||||
}
|
||||
}
|
||||
|
||||
template <typename IntT>
|
||||
__global__ void MaskTable(const IntT* x_indices,
|
||||
const int n,
|
||||
int* index_flags,
|
||||
int* table) {
|
||||
CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
|
||||
int index = x_indices[i];
|
||||
funcs::sparse::SetBits(index, index_flags);
|
||||
table[index] = i;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT, int VecSize>
|
||||
__global__ void MaskCopy(const IntT* mask_indices,
|
||||
const int* index_flags,
|
||||
const int* table,
|
||||
const int n,
|
||||
const int stride,
|
||||
const T* x_values,
|
||||
T* out_values) {
|
||||
using LoadT = AlignedVector<T, VecSize>;
|
||||
using StoreT = AlignedVector<T, VecSize>;
|
||||
CUDA_KERNEL_LOOP_TYPE(i, n, int64_t) {
|
||||
const int mask_index = mask_indices[i];
|
||||
const bool flag = funcs::sparse::TestBits(mask_index, index_flags);
|
||||
if (flag) {
|
||||
int j = table[mask_index];
|
||||
for (int k = 0; k < stride; k += VecSize) {
|
||||
LoadT vec_x;
|
||||
Load<T, VecSize>(x_values + j * stride + k, &vec_x);
|
||||
Store<T, VecSize>(vec_x, out_values + i * stride + k);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename IntT>
|
||||
void MaskHelperCooGPUKernel(const GPUContext& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const DenseTensor& mask_indices,
|
||||
DenseTensor* out) {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
mask_indices.dims().size(),
|
||||
2,
|
||||
common::errors::InvalidArgument("the mask_indices must be 2-D tensor"));
|
||||
|
||||
const int32_t sparse_dim = x.sparse_dim();
|
||||
auto indices_dtype = phi::CppTypeToDataType<IntT>::Type();
|
||||
|
||||
std::vector<IntT> sparse_offsets(sparse_dim);
|
||||
|
||||
DenseTensorMeta x_indices_meta(indices_dtype, {x.nnz()}, DataLayout::NCHW);
|
||||
DenseTensorMeta mask_indices_meta(
|
||||
indices_dtype, {mask_indices.dims()[1]}, DataLayout::NCHW);
|
||||
DenseTensorMeta sparse_offset_meta(
|
||||
indices_dtype, {sparse_dim}, DataLayout::NCHW);
|
||||
|
||||
DenseTensor x_indices = Empty<GPUContext>(dev_ctx, std::move(x_indices_meta));
|
||||
DenseTensor mask_meta_indices =
|
||||
Empty<GPUContext>(dev_ctx, std::move(mask_indices_meta));
|
||||
DenseTensor bound_out =
|
||||
Empty<GPUContext>(dev_ctx, std::move(mask_indices_meta));
|
||||
DenseTensor d_sparse_offsets =
|
||||
Empty<GPUContext>(dev_ctx, std::move(sparse_offset_meta));
|
||||
IntT* x_indices_ptr = x_indices.data<IntT>();
|
||||
IntT* mask_indices_ptr = mask_meta_indices.data<IntT>();
|
||||
IntT* bound_out_ptr = bound_out.data<IntT>();
|
||||
|
||||
// 1. calc the offsets of per dim
|
||||
funcs::sparse::CalcOffsetsPerDim(x.dims(), sparse_dim, sparse_offsets.data());
|
||||
// 2. copy sparse_offsets to device
|
||||
backends::gpu::GpuMemcpyAsync(d_sparse_offsets.data<IntT>(),
|
||||
sparse_offsets.data(),
|
||||
sizeof(IntT) * sparse_dim,
|
||||
gpuMemcpyHostToDevice,
|
||||
dev_ctx.stream());
|
||||
|
||||
// 3. flatten x indices and mask indices
|
||||
auto config =
|
||||
backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_indices.numel(), 1);
|
||||
funcs::sparse::FlattenIndicesKernel<<<config.block_per_grid,
|
||||
config.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
x.indices().data<IntT>(),
|
||||
d_sparse_offsets.data<IntT>(),
|
||||
x_indices.numel(),
|
||||
sparse_dim,
|
||||
x_indices_ptr);
|
||||
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(
|
||||
dev_ctx, mask_meta_indices.numel(), 1);
|
||||
funcs::sparse::FlattenIndicesKernel<<<config.block_per_grid,
|
||||
config.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(
|
||||
mask_indices.data<IntT>(),
|
||||
d_sparse_offsets.data<IntT>(),
|
||||
mask_meta_indices.numel(),
|
||||
sparse_dim,
|
||||
mask_indices_ptr);
|
||||
|
||||
int table_size = 1;
|
||||
auto x_dims = x.dims();
|
||||
for (int i = 0; i < sparse_dim; i++) {
|
||||
table_size *= x_dims[i];
|
||||
}
|
||||
DenseTensor table = Empty<int>(dev_ctx, {table_size});
|
||||
DenseTensor index_flags = Empty<int>(dev_ctx, {(table_size + 31) / 32});
|
||||
backends::gpu::GpuMemsetAsync(index_flags.data<int>(),
|
||||
0,
|
||||
index_flags.numel() * sizeof(int),
|
||||
dev_ctx.stream());
|
||||
const int64_t stride =
|
||||
x.dims().size() == sparse_dim ? 1 : x.values().dims()[1];
|
||||
*out = EmptyLike<T>(dev_ctx, x.values());
|
||||
funcs::SetConstant<GPUContext, T> set_zero;
|
||||
set_zero(dev_ctx, out, static_cast<T>(0));
|
||||
T* out_ptr = out->data<T>();
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(dev_ctx, x_indices.numel(), 1);
|
||||
MaskTable<<<config.block_per_grid,
|
||||
config.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(x_indices_ptr,
|
||||
x_indices.numel(),
|
||||
index_flags.data<int>(),
|
||||
table.data<int>());
|
||||
config = backends::gpu::GetGpuLaunchConfig1D(
|
||||
dev_ctx, mask_meta_indices.numel(), 1);
|
||||
|
||||
const int VecBytes = 16;
|
||||
const int VecSize = VecBytes / sizeof(T);
|
||||
if (stride % VecSize == 0) {
|
||||
MaskCopy<T, IntT, VecSize><<<config.block_per_grid,
|
||||
config.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(mask_indices_ptr,
|
||||
index_flags.data<int>(),
|
||||
table.data<int>(),
|
||||
mask_meta_indices.numel(),
|
||||
stride,
|
||||
x.values().data<T>(),
|
||||
out_ptr);
|
||||
} else {
|
||||
MaskCopy<T, IntT, 1><<<config.block_per_grid,
|
||||
config.thread_per_block,
|
||||
0,
|
||||
dev_ctx.stream()>>>(mask_indices_ptr,
|
||||
index_flags.data<int>(),
|
||||
table.data<int>(),
|
||||
mask_meta_indices.numel(),
|
||||
stride,
|
||||
x.values().data<T>(),
|
||||
out_ptr);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename Context>
|
||||
void MaskHelperCooKernel(const Context& dev_ctx,
|
||||
const SparseCooTensor& x,
|
||||
const DenseTensor& mask_indices,
|
||||
DenseTensor* out) {
|
||||
PD_VISIT_BASE_INTEGRAL_TYPES(
|
||||
x.indices().dtype(), "MaskHelperCooGPUKernel", ([&] {
|
||||
MaskHelperCooGPUKernel<T, data_t>(dev_ctx, x, mask_indices, out);
|
||||
}));
|
||||
}
|
||||
|
||||
} // namespace sparse
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(mask_helper_coo,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::MaskHelperCooKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
uint8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(0).SetDataLayout(phi::DataLayout::SPARSE_COO);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(mask_as_coo,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::MaskAsCooKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_COO);
|
||||
}
|
||||
|
||||
PD_REGISTER_KERNEL(mask_as_csr,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::sparse::MaskAsCsrKernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
uint8_t,
|
||||
int8_t,
|
||||
int16_t,
|
||||
int,
|
||||
int64_t,
|
||||
bool,
|
||||
phi::complex64,
|
||||
phi::complex128) {
|
||||
kernel->InputAt(1).SetDataLayout(phi::DataLayout::SPARSE_CSR);
|
||||
}
|
||||
Reference in New Issue
Block a user