记录一下关于torch的算子优化。主要是根据torch的torch.utils.cpp_extension 去把一些torch上的python变换放到gpu中去并行执行，写好对应的测试用例保证c和python的结果一致。
https://github.com/PrincetonLIPS/numpy-hilbert-curve 中有python版本的希尔伯特编码的实现。目前是期望提速该函数
1 编写cuda核函数
这里的希尔伯特编码是批量进行的，正好对应cuda中的线程，每个线程就可以对一个点进行编码
编写c函数供外部调用，外加pybind进行封装提供python api 供python调用
hilbert_cuda.cu
cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void hilbert_encode_kernel(const uint8_t* locs_uint8, int num_dims, int num_bits, int64_t* output, int N)
{
int tid = blockIdx.x * blockDim.x + threadIdx.x;
if (tid >= N) return;
output[tid] = 0;
__shared__ bool s_gray[256][64];
bool* g = s_gray[threadIdx.x];
const uint8_t* src = locs_uint8 + tid * num_dims * 8;
// 1. 位提取
uint8_t bitpack_mask_rev[8] = {128, 64, 32, 16, 8, 4, 2, 1};
for (int dim = 0; dim < num_dims; dim++) {
// 与 python locs_uint8.flip(-1) 对齐，翻转 8 个字节
uint8_t flipped_bytes[8];
#pragma unroll
for (int i = 0; i < 8; i++) {
// 翻转顺序读取
flipped_bytes[i] = src[dim * 8 + (7 - i)];
}
// 依次从高位到低位提取 8*8=64 bits
bool bits_64[64];
int idx = 0;
#pragma unroll
for (int i = 0; i < 8; i++) {
uint8_t val = flipped_bytes[i];
#pragma unroll
for (int b = 0; b < 8; b++) {
bits_64[idx++] = (val & bitpack_mask_rev[b]) != 0;
}
}
// 与 Python 保持一致，只取最后 num_bits 位
int start_idx = 64 - num_bits;
for (int b = 0; b < num_bits; b++) {
g[dim * num_bits + b] = bits_64[start_idx + b];
}
}
// 2. Skilling变换
for (int bit = 0; bit < num_bits; bit++) {
for (int dim = 0; dim < num_dims; dim++) {
if (bit + 1 >= num_bits) continue;
bool mask = g[dim * num_bits + bit];
// 修改整个bit+1到最后的区间
for (int b = bit + 1; b < num_bits; b++) {
bool orig_0 = g[0 * num_bits + b];
bool orig_dim = g[dim * num_bits + b];
g[0 * num_bits + b] ^= mask;
bool to_flip = (!mask) & (g[0 * num_bits + b] ^ orig_dim);
g[dim * num_bits + b] ^= to_flip;
g[0 * num_bits + b] ^= to_flip;
}
}
}
// 3. Gray转Binary
// 先调整位序为bit-first格式
bool temp[64]; // 临时数组存储转换后的数据
for (int b = 0; b < num_bits; b++) {
for (int dim = 0; dim < num_dims; dim++) {
temp[b * num_dims + dim] = g[dim * num_bits + b];
}
}
// 2. 复制回原数组
for (int i = 0; i < num_bits * num_dims; i++) {
g[i] = temp[i];
}
// 3. Gray to Binary转换
for (int i = 1; i < num_bits * num_dims; i++) {
g[i] = g[i] ^ g[i - 1];
}
// 4.1 模拟padding操作
const int total_bits = 64;
const int extra_bits = total_bits - num_bits * num_dims;
bool padded[64] = {0}; // 初始化为0
// 复制原始数据到padded的后面部分
for(int i = 0; i < num_bits * num_dims; i++) {
padded[extra_bits + i] = g[i];
}
// 4.2 模拟flip操作
bool flipped[64];
#pragma unroll
for(int i = 0; i < 64; i++) {
flipped[i] = padded[63 - i];
}
// 4.3 位打包
int64_t result = 0;
#pragma unroll
for(int i = 0; i < 64; i++) {
if(flipped[i]) {
int64_t bit_value = 1ULL << i;
result |= bit_value;
}
}
output[tid] = result;
}
// 示例 CUDA 内核签名
extern "C" void hilbert_encode_cuda_optimized(torch::Tensor locs_uint8, int num_dims, int num_bits, torch::Tensor output) {
// 获取指针
auto locs_ptr = locs_uint8.data_ptr<uint8_t>();
auto output_ptr = output.data_ptr<int64_t>();
// CUDA核函数调用
int N = locs_uint8.size(0);
dim3 blocks((N + 255) / 256);
dim3 threads(256);
hilbert_encode_kernel<<<blocks, threads>>>(locs_ptr, num_dims, num_bits, output_ptr, N);
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("hilbert_encode_cuda_optimized",
&hilbert_encode_cuda_optimized,
"Optimized Hilbert Encode CUDA");
}2.1 load形式调用
torch.utils.cpp_extension 里的load会运行时根据输入文件，编译这个so库，记得环境变量里要设置好对应的cpp版本等
cpp
# 2) 编译并加载上面的CUDA扩展
cuda_path = os.path.join(os.path.dirname(__file__), "hilbert_cuda.cu")
hilbert_cuda = load(
name="hilbert_cuda",
sources=[cuda_path],
verbose=True
)
def encode_cuda_optimized(locs, num_dims, num_bits):
device = locs.device
N = locs.shape[0]
# 格式化输入
locs_uint8 = locs.contiguous().view(torch.uint8).reshape(N, num_dims, 8)
# 分配输出
output = torch.empty(N, dtype=torch.int64, device=device)
# 执行kernel，并传递output
with torch.cuda.device(device):
hilbert_cuda.hilbert_encode_cuda_optimized(locs_uint8, num_dims, num_bits, output)
torch.cuda.synchronize()
return outputtest.py
python
"""
Hilbert Order
Modified from https://github.com/PrincetonLIPS/numpy-hilbert-curve
Author: Xiaoyang Wu (xiaoyang.wu.cs@gmail.com), Kaixin Xu
Please cite our work if the code is helpful to you.
"""
import torch
from torch.utils.cpp_extension import load
import os
import time
# 2) 编译并加载上面的CUDA扩展
cuda_path = os.path.join(os.path.dirname(__file__), "hilbert_cuda.cu")
hilbert_cuda = load(
name="hilbert_cuda",
sources=[cuda_path],
verbose=True
)
def encode_cuda_optimized(locs, num_dims, num_bits):
device = locs.device
N = locs.shape[0]
# 格式化输入
locs_uint8 = locs.contiguous().view(torch.uint8).reshape(N, num_dims, 8)
# 分配输出
output = torch.empty(N, dtype=torch.int64, device=device)
# 执行kernel，并传递output
with torch.cuda.device(device):
hilbert_cuda.hilbert_encode_cuda_optimized(locs_uint8, num_dims, num_bits, output)
torch.cuda.synchronize()
return output
def right_shift(binary, k=1, axis=-1):
"""Right shift an array of binary values.
Parameters:
-----------
binary: An ndarray of binary values.
k: The number of bits to shift. Default 1.
axis: The axis along which to shift. Default -1.
Returns:
--------
Returns an ndarray with zero prepended and the ends truncated, along
whatever axis was specified."""
# If we're shifting the whole thing, just return zeros.
if binary.shape[axis] <= k:
return torch.zeros_like(binary)
# Determine the padding pattern.
# padding = [(0,0)] * len(binary.shape)
# padding[axis] = (k,0)
# Determine the slicing pattern to eliminate just the last one.
slicing = [slice(None)] * len(binary.shape)
slicing[axis] = slice(None, -k)
shifted = torch.nn.functional.pad(
binary[tuple(slicing)], (k, 0), mode="constant", value=0
)
return shifted
def gray2binary_old(gray, axis=-1):
"""Convert an array of Gray codes back into binary values.
Parameters:
-----------
gray: An ndarray of gray codes.
axis: The axis along which to perform Gray decoding. Default=-1.
Returns:
--------
Returns an ndarray of binary values.
"""
# Loop the log2(bits) number of times necessary, with shift and xor.
shift = 2 ** (torch.Tensor([gray.shape[axis]]).log2().ceil().int() - 1)
while shift > 0:
gray = torch.logical_xor(gray, right_shift(gray, shift))
shift = torch.div(shift, 2, rounding_mode="floor")
return gray
def encode(locs, num_dims, num_bits):
"""Decode an array of locations in a hypercube into a Hilbert integer.
This is a vectorized-ish version of the Hilbert curve implementation by John
Skilling as described in:
Skilling, J. (2004, April). Programming the Hilbert curve. In AIP Conference
Proceedings (Vol. 707, No. 1, pp. 381-387). American Institute of Physics.
Params:
-------
locs - An ndarray of locations in a hypercube of num_dims dimensions, in
which each dimension runs from 0 to 2**num_bits-1. The shape can
be arbitrary, as long as the last dimension of the same has size
num_dims.
num_dims - The dimensionality of the hypercube. Integer.
num_bits - The number of bits for each dimension. Integer.
Returns:
--------
The output is an ndarray of uint64 integers with the same shape as the
input, excluding the last dimension, which needs to be num_dims.
"""
# Keep around the original shape for later.
orig_shape = locs.shape
bitpack_mask = 1 << torch.arange(0, 8).to(locs.device)
bitpack_mask_rev = bitpack_mask.flip(-1)
if orig_shape[-1] != num_dims:
raise ValueError(
"""
The shape of locs was surprising in that the last dimension was of size
%d, but num_dims=%d. These need to be equal.
"""
% (orig_shape[-1], num_dims)
)
if num_dims * num_bits > 63:
raise ValueError(
"""
num_dims=%d and num_bits=%d for %d bits total, which can't be encoded
into a int64. Are you sure you need that many points on your Hilbert
curve?
"""
% (num_dims, num_bits, num_dims * num_bits)
)
# Treat the location integers as 64-bit unsigned and then split them up into
# a sequence of uint8s. Preserve the association by dimension.
locs_uint8 = locs.long().view(torch.uint8).reshape((-1, num_dims, 8)).flip(-1)
# Now turn these into bits and truncate to num_bits.
gray = (
locs_uint8.unsqueeze(-1)
.bitwise_and(bitpack_mask_rev)
.ne(0)
.byte()
.flatten(-2, -1)[..., -num_bits:]
)
# Run the decoding process the other way.
# Iterate forwards through the bits.
for bit in range(0, num_bits):
# Iterate forwards through the dimensions.
for dim in range(0, num_dims):
# Identify which ones have this bit active.
mask = gray[:, dim, bit]
# Where this bit is on, invert the 0 dimension for lower bits.
gray[:, 0, bit + 1 :] = torch.logical_xor(
gray[:, 0, bit + 1 :], mask[:, None]
)
# Where the bit is off, exchange the lower bits with the 0 dimension.
to_flip = torch.logical_and(
torch.logical_not(mask[:, None]).repeat(1, gray.shape[2] - bit - 1),
torch.logical_xor(gray[:, 0, bit + 1 :], gray[:, dim, bit + 1 :]),
)
gray[:, dim, bit + 1 :] = torch.logical_xor(
gray[:, dim, bit + 1 :], to_flip
)
gray[:, 0, bit + 1 :] = torch.logical_xor(gray[:, 0, bit + 1 :], to_flip)
# Now flatten out.
gray = gray.swapaxes(1, 2).reshape((-1, num_bits * num_dims))
# Convert Gray back to binary.
hh_bin = gray2binary_old(gray)
# Pad back out to 64 bits.
extra_dims = 64 - num_bits * num_dims
padded = torch.nn.functional.pad(hh_bin, (extra_dims, 0), "constant", 0)
# Convert binary values into uint8s.
hh_uint8 = (
(padded.flip(-1).reshape((-1, 8, 8)) * bitpack_mask)
.sum(2)
.squeeze()
.type(torch.uint8)
)
# Convert uint8s into uint64s.
hh_uint64 = hh_uint8.view(torch.int64).squeeze()
return hh_uint64
def encode_optimized(locs, num_dims, num_bits):
"""
与 encode 数值等价、性能更优的 Hilbert 曲线编码，支持 GPU 并行。
"""
shp = locs.shape
if num_dims * num_bits > 63:
raise ValueError("Cannot encode more than 63 bits in total.")
device = locs.device
# print("\n=== Python Version Debug ===")
# print(f"Input shape: {locs.shape}")
# print(f"First sample raw data:")
# for dim in range(num_dims):
# print(f"dim {dim}: {locs[1,dim]:08b} ({locs[1,dim]})") # 添加原始二进制表示
# 与原版一致的步骤：先按字节拆分并翻转
locs_uint8 = locs.view(torch.uint8).reshape(-1, num_dims, 8).flip(-1)
# print("\n=== Bit Extraction Process ===")
# print("1. After reshape & flip:")
# print(f"locs_uint8[0]:\n{locs_uint8[1]}")
bitpack_mask = 1 << torch.arange(0, 8, device=device)
bitpack_mask_rev = bitpack_mask.flip(-1)
# print(f"bitpack_mask_rev: {bitpack_mask_rev}")
# 提取高位到低位（与原 encode 的做法一致）
gray = (
locs_uint8.unsqueeze(-1)
.bitwise_and(bitpack_mask_rev)
.ne(0)
.byte()
.flatten(-2, -1)[..., -num_bits:]
.reshape(-1, num_dims, num_bits)
)
# print("\n=== After bit extraction ===")
# print(f"First sample bits by dimension:")
# for dim in range(num_dims):
# bits = ""
# for b in range(num_bits):
# bits += str(int(gray[1,dim,b]))
# print(f"dim {dim}: {bits}")
# Skilling 旋转/翻转
for bit in range(num_bits):
for dim in range(num_dims):
mask = gray[:, dim, bit]
if bit + 1 < num_bits:
gray[:, 0, bit+1:] ^= mask.unsqueeze(-1)
to_flip = (~mask.unsqueeze(-1)) & (gray[:, 0, bit+1:] ^ gray[:, dim, bit+1:])
gray[:, dim, bit+1:] ^= to_flip
gray[:, 0, bit+1:] ^= to_flip
# print("\n=== After Skilling Transform ===")
# print(f"First sample bits after Skilling:")
# for dim in range(num_dims):
# bits = ""
# for b in range(num_bits):
# bits += str(int(gray[1,dim,b]))
# print(f"dim {dim}: {bits}")
# Gray→binary
gview = gray.swapaxes(1, 2).reshape(-1, num_bits * num_dims)
# print("\n=== After swapaxes ===")
# print(f"Bit order after swapaxes:")
# print(gview[1])
binary = torch.cumsum(gview, dim=-1) % 2
# print("\n=== After Gray to Binary ===")
# print(f"First sample binary representation:")
# print(f"binary[0]: {binary[1].tolist()}")
binary = binary.long() # 转换为整数类型
# 填充到 64 位并打包
extra_bits = 64 - num_bits * num_dims
padded = torch.nn.functional.pad(binary, (extra_bits, 0)).long()
# 使用整数类型的 bitmask
bitmask = 1 << torch.arange(64, device=device).long()
hilbert_indices = (padded.flip(-1) * bitmask).sum(dim=-1).long()
# print("\n=== Final Result ===")
# print(f"Final packed result: {hilbert_indices[1]}")
return hilbert_indices.reshape(shp[:-1])
# 测试用例
torch.manual_seed(41) #可修改
# 假设我们有 3 个维度
num_dims = 3 # 3维
import random
npoints = 10
total_time_encode = 0
total_time_encode_new = 0
num_runs = 10 # 可根据需要调整测试次数
for i in range(0, num_runs):
num_bits = random.randint(3, 20)
#num_bits = 3
torch.manual_seed(i)
locs1 = torch.randint(0, 2**num_bits, (npoints, num_dims), dtype=torch.int64)
locs2 = locs1.to('cuda')
num_bits = depth = int(torch.max(locs1)).bit_length()
print(num_bits)
# 测试 encode
start_time = time.time()
encoded_values_encode = encode_optimized(locs1, num_dims, num_bits)
end_time = time.time()
total_time_encode += (end_time - start_time)
# 测试 encode_new
torch.cuda.synchronize()
start_time = time.time()
encoded_values_encode_new = encode_cuda_optimized(locs2, num_dims, num_bits)
end_time = time.time()
total_time_encode_new += (end_time - start_time)
encoded_values_encode_new_cpu = encoded_values_encode_new.cpu()
# 比较 encode 和 encode_new 的结果
if torch.equal(encoded_values_encode, encoded_values_encode_new_cpu):
print("encode 和 encode_new 生成的编码结果完全相同。")
else:
print("encode 和 encode_new 生成的编码结果存在差异。")
# 计算不匹配的数量
mismatches = (encoded_values_encode != encoded_values_encode_new_cpu).sum().item()
print(f"不匹配的元素数量: {mismatches}")
if mismatches > 0:
mismatched_indices = (encoded_values_encode != encoded_values_encode_new_cpu).nonzero(as_tuple=True)[0]
num_to_print = min(10, mismatches)
print(f"前 {num_to_print} 个不匹配的示例:")
for i in mismatched_indices[:num_to_print]:
print(f"索引 {i.item()}: encode = {encoded_values_encode[i].item()}, encode_new = {encoded_values_encode_new[i].item()}")
#分别打印一下encode以及encode_optimized测试耗时
print("encode 耗时: ", total_time_encode/num_runs)
print("encode_optimized 耗时: ", total_time_encode_new/num_runs)
# Encoded Values:
# tensor([12502877, 15055001, 12711671, ..., 10978447, 5159808, 12237478])2.2 import形式调用
setup.py
cpp
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os
setup(
name='hilbert_cuda', # 扩展模块的名称
ext_modules=[
CUDAExtension(
name='hilbert_cuda', # 与 Python 中加载的模块名一致
sources=['hilbert_cuda.cu'], # 源文件列表
extra_compile_args={
'cxx': ['-O3'], # C++ 编译器优化选项
'nvcc': [
'-O3', # CUDA 编译器优化选项
'--use_fast_math', # 使用快速数学库
'-gencode=arch=compute_89,code=sm_89', # 针对 RTX 4090 的架构
'-std=c++17', # 使用 C++17 标准
'--expt-relaxed-constexpr', # 允许更宽松的常量表达式
],
},
include_dirs=[], # 如有需要，可添加额外的包含目录
),
],
cmdclass={
'build_ext': BuildExtension # 指定构建扩展的命令类
}
)build
执行完毕后对应目录下会生成 对应的hilbert_cuda.cpython-38-x86_64-linux-gnu.so，python调用只看名字，是c的so也没问题
bash
[root@218c3c31e907 6]# python3 setup.py build
running build
running build_ext
/usr/local/lib/python3.8/site-packages/torch/utils/cpp_extension.py:398: UserWarning: There are no /usr/local/gcc-11.4/bin/g++ version bounds defined for CUDA version 11.8
warnings.warn(f'There are no {compiler_name} version bounds defined for CUDA version {cuda_str_version}')
building 'hilbert_cuda' extension
Emitting ninja build file /zyx/6/build/temp.linux-x86_64-3.8/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
/usr/local/gcc-11.4/bin/g++ -pthread -shared /zyx/6/build/temp.linux-x86_64-3.8/hilbert_cuda.o -L/usr/local/lib/python3.8/site-packages/torch/lib -L/usr/local/cuda/lib64 -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda -o build/lib.linux-x86_64-3.8/hilbert_cuda.cpython-38-x86_64-linux-gnu.so此时在对应有so的目录下就不会报错了
python
import torch
import hilbert_cuda install
执行完毕后会把包安装到对应的python的包路径下，就可以直接调用了
bash
[root@218c3c31e907 hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg]# pip show hilbert_cuda
Name: hilbert-cuda
Version: 0.0.0
Summary: UNKNOWN
Home-page: UNKNOWN
Author:
Author-email:
License: UNKNOWN
Location: /usr/local/lib/python3.8/site-packages/hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg
Requires:
Required-by:python3 setup.py install
bash
[root@218c3c31e907 6]# python3 setup.py install
running install
running bdist_egg
running egg_info
creating hilbert_cuda.egg-info
writing hilbert_cuda.egg-info/PKG-INFO
writing dependency_links to hilbert_cuda.egg-info/dependency_links.txt
writing top-level names to hilbert_cuda.egg-info/top_level.txt
writing manifest file 'hilbert_cuda.egg-info/SOURCES.txt'
listing git files failed - pretending there aren't any
reading manifest file 'hilbert_cuda.egg-info/SOURCES.txt'
writing manifest file 'hilbert_cuda.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_ext
/usr/local/lib/python3.8/site-packages/torch/utils/cpp_extension.py:398: UserWarning: There are no /usr/local/gcc-11.4/bin/g++ version bounds defined for CUDA version 11.8
warnings.warn(f'There are no {compiler_name} version bounds defined for CUDA version {cuda_str_version}')
building 'hilbert_cuda' extension
Emitting ninja build file /zyx/6/build/temp.linux-x86_64-3.8/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
/usr/local/gcc-11.4/bin/g++ -pthread -shared /zyx/6/build/temp.linux-x86_64-3.8/hilbert_cuda.o -L/usr/local/lib/python3.8/site-packages/torch/lib -L/usr/local/cuda/lib64 -lc10 -ltorch -ltorch_cpu -ltorch_python -lcudart -lc10_cuda -ltorch_cuda -o build/lib.linux-x86_64-3.8/hilbert_cuda.cpython-38-x86_64-linux-gnu.so
creating build/bdist.linux-x86_64
creating build/bdist.linux-x86_64/egg
copying build/lib.linux-x86_64-3.8/hilbert_cuda.cpython-38-x86_64-linux-gnu.so -> build/bdist.linux-x86_64/egg
creating stub loader for hilbert_cuda.cpython-38-x86_64-linux-gnu.so
byte-compiling build/bdist.linux-x86_64/egg/hilbert_cuda.py to hilbert_cuda.cpython-38.pyc
creating build/bdist.linux-x86_64/egg/EGG-INFO
copying hilbert_cuda.egg-info/PKG-INFO -> build/bdist.linux-x86_64/egg/EGG-INFO
copying hilbert_cuda.egg-info/SOURCES.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying hilbert_cuda.egg-info/dependency_links.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
copying hilbert_cuda.egg-info/top_level.txt -> build/bdist.linux-x86_64/egg/EGG-INFO
writing build/bdist.linux-x86_64/egg/EGG-INFO/native_libs.txt
zip_safe flag not set; analyzing archive contents...
__pycache__.hilbert_cuda.cpython-38: module references __file__
creating dist
creating 'dist/hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg' and adding 'build/bdist.linux-x86_64/egg' to it
removing 'build/bdist.linux-x86_64/egg' (and everything under it)
Processing hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg
creating /usr/local/lib/python3.8/site-packages/hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg
Extracting hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg to /usr/local/lib/python3.8/site-packages
Adding hilbert-cuda 0.0.0 to easy-install.pth file
Installed /usr/local/lib/python3.8/site-packages/hilbert_cuda-0.0.0-py3.8-linux-x86_64.egg
Processing dependencies for hilbert-cuda==0.0.0
Finished processing dependencies for hilbert-cuda==0.0.0test_setup.py
python
import torch
import hilbert_cuda # 导入编译好的扩展模块
import os
import time
import random
def encode_cuda_optimized(locs, num_dims, num_bits):
device = locs.device
N = locs.shape[0]
# 格式化输入
locs_uint8 = locs.contiguous().view(torch.uint8).reshape(N, num_dims, 8)
# 分配输出
output = torch.empty(N, dtype=torch.int64, device=device)
# 执行 kernel，并传递 output
with torch.cuda.device(device):
hilbert_cuda.hilbert_encode_cuda_optimized(locs_uint8, num_dims, num_bits, output)
torch.cuda.synchronize()
return output
def right_shift(binary, k=1, axis=-1):
"""Right shift an array of binary values."""
if binary.shape[axis] <= k:
return torch.zeros_like(binary)
slicing = [slice(None)] * len(binary.shape)
slicing[axis] = slice(None, -k)
shifted = torch.nn.functional.pad(
binary[tuple(slicing)], (k, 0), mode="constant", value=0
)
return shifted
def gray2binary_old(gray, axis=-1):
"""Convert an array of Gray codes back into binary values."""
shift = 2 ** (torch.Tensor([gray.shape[axis]]).log2().ceil().int() - 1)
while shift > 0:
gray = torch.logical_xor(gray, right_shift(gray, shift))
shift = torch.div(shift, 2, rounding_mode="floor")
return gray
def encode(locs, num_dims, num_bits):
"""CPU 实现的 Hilbert 编码"""
orig_shape = locs.shape
bitpack_mask = 1 << torch.arange(0, 8, device=locs.device)
bitpack_mask_rev = bitpack_mask.flip(-1)
if orig_shape[-1] != num_dims:
raise ValueError(
f"The shape of locs was surprising in that the last dimension was of size {orig_shape[-1]}, but num_dims={num_dims}. These need to be equal."
)
if num_dims * num_bits > 63:
raise ValueError(
f"num_dims={num_dims} and num_bits={num_bits} for {num_dims * num_bits} bits total, which can't be encoded into a int64. Are you sure you need that many points on your Hilbert curve?"
)
locs_uint8 = locs.long().view(torch.uint8).reshape((-1, num_dims, 8)).flip(-1)
gray = (
locs_uint8.unsqueeze(-1)
.bitwise_and(bitpack_mask_rev)
.ne(0)
.byte()
.flatten(-2, -1)[..., -num_bits:]
)
for bit in range(0, num_bits):
for dim in range(0, num_dims):
mask = gray[:, dim, bit]
gray[:, 0, bit + 1 :] = torch.logical_xor(
gray[:, 0, bit + 1 :], mask[:, None]
)
to_flip = torch.logical_and(
torch.logical_not(mask[:, None]).repeat(1, gray.shape[2] - bit - 1),
torch.logical_xor(gray[:, 0, bit + 1 :], gray[:, dim, bit + 1 :]),
)
gray[:, dim, bit + 1 :] = torch.logical_xor(
gray[:, dim, bit + 1 :], to_flip
)
gray[:, 0, bit + 1 :] = torch.logical_xor(
gray[:, 0, bit + 1 :], to_flip
)
gray = gray.swapaxes(1, 2).reshape(-1, num_bits * num_dims)
hh_bin = gray2binary_old(gray)
extra_dims = 64 - num_bits * num_dims
padded = torch.nn.functional.pad(hh_bin, (extra_dims, 0), "constant", 0)
hh_uint8 = (
(padded.flip(-1).reshape((-1, 8, 8)) * bitpack_mask)
.sum(2)
.squeeze()
.type(torch.uint8)
)
hh_uint64 = hh_uint8.view(torch.int64).squeeze()
return hh_uint64.reshape(orig_shape[:-1])
def encode_optimized(locs, num_dims, num_bits):
"""
与 encode 数值等价、性能更优的 Hilbert 曲线编码，支持 GPU 并行。
"""
shp = locs.shape
if num_dims * num_bits > 63:
raise ValueError("Cannot encode more than 63 bits in total.")
device = locs.device
locs_uint8 = locs.view(torch.uint8).reshape(-1, num_dims, 8)
bitpack_mask = 1 << torch.arange(0, 8, device=device)
bitpack_mask_rev = bitpack_mask.flip(-1)
gray = (
locs_uint8.unsqueeze(-1)
.bitwise_and(bitpack_mask_rev)
.ne(0)
.byte()
.flatten(-2, -1)[..., -num_bits:]
.reshape(-1, num_dims, num_bits)
)
# Skilling 变换
for bit in range(num_bits):
for dim in range(num_dims):
mask = gray[:, dim, bit]
if bit + 1 < num_bits:
gray[:, 0, bit+1:] ^= mask.unsqueeze(-1)
to_flip = (~mask.unsqueeze(-1)) & (gray[:, 0, bit+1:] ^ gray[:, dim, bit+1:])
gray[:, dim, bit+1:] ^= to_flip
gray[:, 0, bit+1:] ^= to_flip
# Gray 转 Binary
gview = gray.swapaxes(1, 2).reshape(-1, num_bits * num_dims)
binary = torch.cumsum(gview, dim=-1) % 2
binary = binary.long()
# 填充到 64 位并打包
extra_bits = 64 - num_bits * num_dims
padded = torch.nn.functional.pad(binary, (extra_bits, 0)).long()
bitmask = 1 << torch.arange(64, device=device).long()
hilbert_indices = (padded.flip(-1) * bitmask).sum(dim=-1).long()
return hilbert_indices.reshape(shp[:-1])
# 测试用例
def main():
torch.manual_seed(41) # 可修改
# 假设我们有 3 个维度
num_dims = 3 # 3维
npoints = 10
total_time_encode = 0
total_time_encode_new = 0
num_runs = 10 # 可根据需要调整测试次数
for i in range(0, num_runs):
num_bits = random.randint(3, 20)
torch.manual_seed(i)
locs1 = torch.randint(0, 2**num_bits, (npoints, num_dims), dtype=torch.int64)
locs2 = locs1.to('cuda')
num_bits = depth = int(torch.max(locs1)).bit_length()
print(f"Run {i+1}: num_bits = {num_bits}")
# 测试 encode
start_time = time.time()
encoded_values_encode = encode(locs1, num_dims, num_bits)
end_time = time.time()
total_time_encode += (end_time - start_time)
# 测试 encode_optimized
start_time = time.time()
encoded_values_encode_new = encode_cuda_optimized(locs2, num_dims, num_bits)
end_time = time.time()
total_time_encode_new += (end_time - start_time)
encoded_values_encode_new_cpu = encoded_values_encode_new.cpu()
# 比较 encode 和 encode_new 的结果
if torch.equal(encoded_values_encode, encoded_values_encode_new_cpu):
print("encode 和 encode_new 生成的编码结果完全相同。")
else:
print("encode 和 encode_new 生成的编码结果存在差异。")
# 计算不匹配的数量
mismatches = (encoded_values_encode != encoded_values_encode_new_cpu).sum().item()
print(f"不匹配的元素数量: {mismatches}")
if mismatches > 0:
mismatched_indices = (encoded_values_encode != encoded_values_encode_new_cpu).nonzero(as_tuple=True)[0]
num_to_print = min(10, mismatches)
print(f"前 {num_to_print} 个不匹配的示例:")
for idx in mismatched_indices[:num_to_print]:
print(f"索引 {idx.item()}: encode = {encoded_values_encode[idx].item()}, encode_new = {encoded_values_encode_new_cpu[idx].item()}")
# 分别打印 encode 以及 encode_optimized 测试耗时
print("encode 耗时: ", total_time_encode / num_runs)
print("encode_optimized 耗时: ", total_time_encode_new / num_runs)
if __name__ == "__main__":
main()