forked from ccf-ai-infra/GPUCodeForces
100 lines
2.5 KiB
Python
100 lines
2.5 KiB
Python
import torch
|
|
from torch.utils.cpp_extension import load_inline
|
|
|
|
has_element_source = """
|
|
#include <torch/extension.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
__global__ void has_element_vec4_kernel(
|
|
const int* __restrict__ indices,
|
|
float* __restrict__ output,
|
|
int n)
|
|
{
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
|
int vec_idx = idx * 4;
|
|
|
|
// 1. Vectorized Path
|
|
if (vec_idx + 4 <= n) {
|
|
// Load int4
|
|
const int4* in_ptr = reinterpret_cast<const int4*>(indices);
|
|
int4 in_val = in_ptr[idx];
|
|
|
|
float4 out_val;
|
|
|
|
out_val.x = (in_val.x != -1) ? 1.0f : 0.0f;
|
|
out_val.y = (in_val.y != -1) ? 1.0f : 0.0f;
|
|
out_val.z = (in_val.z != -1) ? 1.0f : 0.0f;
|
|
out_val.w = (in_val.w != -1) ? 1.0f : 0.0f;
|
|
|
|
float4* out_ptr = reinterpret_cast<float4*>(output);
|
|
out_ptr[idx] = out_val;
|
|
}
|
|
}
|
|
|
|
__global__ void has_element_scalar_kernel(
|
|
const int* __restrict__ indices,
|
|
float* __restrict__ output,
|
|
int n,
|
|
int offset)
|
|
{
|
|
int idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
|
|
|
|
if (idx < n) {
|
|
int val = indices[idx];
|
|
output[idx] = (val != -1) ? 1.0f : 0.0f;
|
|
}
|
|
}
|
|
|
|
torch::Tensor has_element_cuda(torch::Tensor indices) {
|
|
int n = indices.numel();
|
|
|
|
auto output = torch::empty_like(indices, torch::dtype(torch::kFloat32));
|
|
|
|
|
|
int vec_elements = n / 4;
|
|
if (vec_elements > 0) {
|
|
const int block = 256;
|
|
const int grid = (vec_elements + block - 1) / block;
|
|
|
|
has_element_vec4_kernel<<<grid, block>>>(
|
|
indices.data_ptr<int>(),
|
|
output.data_ptr<float>(),
|
|
n
|
|
);
|
|
}
|
|
|
|
|
|
int remainder = n % 4;
|
|
if (remainder > 0) {
|
|
int offset = vec_elements * 4;
|
|
has_element_scalar_kernel<<<1, 32>>>(
|
|
indices.data_ptr<int>(),
|
|
output.data_ptr<float>(),
|
|
n,
|
|
offset
|
|
);
|
|
}
|
|
|
|
return output;
|
|
}
|
|
"""
|
|
|
|
cpp_source = "torch::Tensor has_element_cuda(torch::Tensor indices);"
|
|
|
|
has_element_module = load_inline(
|
|
name="optional_has_element_extension",
|
|
cpp_sources=cpp_source,
|
|
cuda_sources=has_element_source,
|
|
functions=["has_element_cuda"],
|
|
verbose=True,
|
|
with_cuda=True
|
|
)
|
|
|
|
class ModelNew(torch.nn.Module):
|
|
def __init__(self):
|
|
super(ModelNew, self).__init__()
|
|
self.cuda_op = has_element_module
|
|
|
|
def forward(self, indices):
|
|
|
|
return self.cuda_op.has_element_cuda(indices.int().contiguous()) |