diff --git a/.gitignore b/.gitignore index bbd86a2..434ea17 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ data exp +build +__pycache__ diff --git a/README.md b/README.md index 809dcd0..e30a4b9 100644 --- a/README.md +++ b/README.md @@ -106,6 +106,13 @@ University of Zurich +- Optional: Install CUDA kernels for faster point cloud serialization in attention layers (automatic fallback to the original PyTorch implementation). Modify the ```all_cuda_archs``` in ```libs/serialization/setup.py``` (see PointROPE above). + ```shell + cd libs/serialization + python setup.py install + cd ../.. + ``` + - Additional requirements. The requirements below are optional, and only required for evaluator and PointGroup instance segmentation. * For evaluator: ``` @@ -129,6 +136,7 @@ Your_project/ └── .../ └── LitePT/ ├── litept/ + ├── libs/serialization └── libs/pointrope ``` We also provide a ```demo_use.py``` script that illustrates how to use the standalone LitePT with an example point cloud input. diff --git a/litept/serialization/__init__.py b/libs/serialization/__init__.py similarity index 100% rename from litept/serialization/__init__.py rename to libs/serialization/__init__.py diff --git a/libs/serialization/cuda/bindings.cpp b/libs/serialization/cuda/bindings.cpp new file mode 100644 index 0000000..073122d --- /dev/null +++ b/libs/serialization/cuda/bindings.cpp @@ -0,0 +1,12 @@ +#include + +torch::Tensor hilbert_encode_cuda(torch::Tensor coords, int num_bits); + +torch::Tensor morton_encode_cuda(torch::Tensor coords); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("hilbert_encode", &hilbert_encode_cuda, "Hilbert encode (CUDA)"); + + // morton (or Z-order) encoding + m.def("morton_encode", &morton_encode_cuda, "Morton encode (CUDA)"); +} \ No newline at end of file diff --git a/libs/serialization/cuda/hilbert.cu b/libs/serialization/cuda/hilbert.cu new file mode 100644 index 0000000..d8e852d --- /dev/null +++ b/libs/serialization/cuda/hilbert.cu @@ -0,0 +1,102 @@ +#include +#include +#include + +__device__ inline uint64_t hilbert3d_single(uint32_t x, uint32_t y, uint32_t z, int num_bits) { + uint32_t X = x; + uint32_t Y = y; + uint32_t Z = z; + + // --- Skilling transform --- + uint32_t Q = 1u << (num_bits - 1); + while (Q > 1) { + uint32_t P = Q - 1; + + // Dimension 0 (X) + if (X & Q) { + X ^= P; + } + + // Dimension 1 (Y) + if (Y & Q) { + X ^= P; // Invert X's lower bits + } else { + // Swap lower bits of X and Y + uint32_t t = (X ^ Y) & P; + X ^= t; + Y ^= t; + } + + // Dimension 2 (Z) + if (Z & Q) { + X ^= P; // Invert X's lower bits + } else { + // Swap lower bits of X and Z + uint32_t t = (X ^ Z) & P; + X ^= t; + Z ^= t; + } + + Q >>= 1; + } + + // --- Interleave bits --- + // (This produces the Gray-coded Hilbert integer) + uint64_t h = 0; + for (int i = num_bits - 1; i >= 0; --i) { + h <<= 3; + h |= ((uint64_t)((X >> i) & 1) << 2) | + ((uint64_t)((Y >> i) & 1) << 1) | + ((uint64_t)((Z >> i) & 1)); + } + + // --- Gray to Binary --- + // Skilling's outputs a Gray-coded index. We decode it here on the full interleaved integer. + h ^= (h >> 1); + h ^= (h >> 2); + h ^= (h >> 4); + h ^= (h >> 8); + h ^= (h >> 16); + h ^= (h >> 32); + + return h; +} + + +__global__ void hilbert_kernel( + const uint32_t* __restrict__ coords, + uint64_t* __restrict__ out, + int N, + int num_bits +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N) return; + + uint32_t x = coords[3 * idx + 0]; + uint32_t y = coords[3 * idx + 1]; + uint32_t z = coords[3 * idx + 2]; + + out[idx] = hilbert3d_single(x, y, z, num_bits); +} + +torch::Tensor hilbert_encode_cuda(torch::Tensor coords, int num_bits) { + TORCH_CHECK(coords.is_cuda(), "coords must be CUDA"); + TORCH_CHECK(coords.size(1) == 3, "coords must be (N, 3)"); + + int N = coords.size(0); + + auto coords_u32 = coords.to(torch::kUInt32).contiguous(); + auto out = torch::empty({N}, torch::dtype(torch::kUInt64).device(coords.device())); + + int threads = 256; + int blocks = (N + threads - 1) / threads; + + hilbert_kernel<<>>( + coords_u32.data_ptr(), + out.data_ptr(), + N, + num_bits + ); + + return out; +} \ No newline at end of file diff --git a/libs/serialization/cuda/morton.cu b/libs/serialization/cuda/morton.cu new file mode 100644 index 0000000..4f887e6 --- /dev/null +++ b/libs/serialization/cuda/morton.cu @@ -0,0 +1,55 @@ +#include +#include +#include + +__device__ inline uint64_t split_by_3(uint32_t x) { + uint64_t v = x & 0x1fffff; // up to 21 bits + + v = (v | (v << 32)) & 0x1f00000000ffff; + v = (v | (v << 16)) & 0x1f0000ff0000ff; + v = (v | (v << 8)) & 0x100f00f00f00f00f; + v = (v | (v << 4)) & 0x10c30c30c30c30c3; + v = (v | (v << 2)) & 0x1249249249249249; + + return v; +} + +__device__ inline uint64_t morton3d_single(uint32_t x, uint32_t y, uint32_t z) { + return (split_by_3(x) << 2) | (split_by_3(y) << 1) | split_by_3(z); +} + +__global__ void morton_kernel( + const uint32_t* __restrict__ coords, + uint64_t* __restrict__ out, + int N +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N) return; + + uint32_t x = coords[3 * idx + 0]; + uint32_t y = coords[3 * idx + 1]; + uint32_t z = coords[3 * idx + 2]; + + out[idx] = morton3d_single(x, y, z); +} + +torch::Tensor morton_encode_cuda(torch::Tensor coords) { + TORCH_CHECK(coords.is_cuda(), "coords must be CUDA"); + TORCH_CHECK(coords.size(1) == 3, "coords must be (N,3)"); + + int N = coords.size(0); + + auto coords_u32 = coords.to(torch::kUInt32).contiguous(); + auto out = torch::empty({N}, torch::dtype(torch::kUInt64).device(coords.device())); + + int threads = 256; + int blocks = (N + threads - 1) / threads; + + morton_kernel<<>>( + coords_u32.data_ptr(), + out.data_ptr(), + N + ); + + return out; +} \ No newline at end of file diff --git a/litept/serialization/default.py b/libs/serialization/default.py similarity index 50% rename from litept/serialization/default.py rename to libs/serialization/default.py index 233051e..70fc712 100644 --- a/litept/serialization/default.py +++ b/libs/serialization/default.py @@ -4,20 +4,27 @@ from .hilbert import encode as hilbert_encode_ from .hilbert import decode as hilbert_decode_ -import numpy as np -from colorhash import ColorHash +try: + import serialization_cuda # run `python setup.py install` + # print("Has CUDA-accelerated serialization.") + HAS_CUDA_EXT = True +except ImportError: + # print("No CUDA-accelerated serialization.") + HAS_CUDA_EXT = False +# import numpy as np +# from colorhash import ColorHash -def int_to_plotly_rgb(x): - """Convert 1D torch.Tensor of int into plotly-friendly RGB format. - This operation is deterministic on the int values. - """ - assert isinstance(x, torch.Tensor) - assert x.dim() == 1 - assert not x.is_floating_point() - x = x.cpu().long().numpy() - palette = np.array([ColorHash(i).rgb for i in range(x.max() + 1)]) - return palette[x] +# def int_to_plotly_rgb(x): +# """Convert 1D torch.Tensor of int into plotly-friendly RGB format. +# This operation is deterministic on the int values. +# """ +# assert isinstance(x, torch.Tensor) +# assert x.dim() == 1 +# assert not x.is_floating_point() +# x = x.cpu().long().numpy() +# palette = np.array([ColorHash(i).rgb for i in range(x.max() + 1)]) +# return palette[x] @torch.inference_mode() def encode(grid_coord, batch=None, depth=16, order="z"): @@ -59,10 +66,19 @@ def decode(code, depth=16, order="z"): def z_order_encode(grid_coord: torch.Tensor, depth: int = 16): - x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long() - # we block the support to batch, maintain batched code in Point class - code = z_order_encode_(x, y, z, b=None, depth=depth) - return code + if HAS_CUDA_EXT: + code_cuda = serialization_cuda.morton_encode(grid_coord).long() + # basic differential testing of the CUDA implementation + # x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long() + # # we block the support to batch, maintain batched code in Point class + # code_torch = z_order_encode_(x, y, z, b=None, depth=depth) + # print("morton equal", torch.all(code_cuda == code_torch)) + return code_cuda + else: + x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long() + # we block the support to batch, maintain batched code in Point class + code = z_order_encode_(x, y, z, b=None, depth=depth) + return code def z_order_decode(code: torch.Tensor, depth): @@ -71,8 +87,16 @@ def z_order_decode(code: torch.Tensor, depth): return grid_coord + def hilbert_encode(grid_coord: torch.Tensor, depth: int = 16): - return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth) + if HAS_CUDA_EXT: + code_cuda = serialization_cuda.hilbert_encode(grid_coord.to(torch.uint32).contiguous(), depth).long() + # basic differential testing of the CUDA implementation + # code_torch = hilbert_encode_(grid_coord, num_dims=3, num_bits=depth) + # print("hilbert equal", torch.all(code_cuda == code_torch)) + return code_cuda + else: + return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth) def hilbert_decode(code: torch.Tensor, depth: int = 16): diff --git a/litept/serialization/hilbert.py b/libs/serialization/hilbert.py similarity index 100% rename from litept/serialization/hilbert.py rename to libs/serialization/hilbert.py diff --git a/libs/serialization/setup.py b/libs/serialization/setup.py new file mode 100644 index 0000000..bccfb45 --- /dev/null +++ b/libs/serialization/setup.py @@ -0,0 +1,32 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +# compile for all possible CUDA architectures +# all_cuda_archs = cuda.get_gencode_flags().replace('compute=','arch=').split() +# alternatively, you can list cuda archs that you want, eg: +# check https://developer.nvidia.com/cuda-gpus to find your arch +all_cuda_archs = [ + # '-gencode', 'arch=compute_90,code=sm_90', + # '-gencode', 'arch=compute_75,code=sm_75', + # '-gencode', 'arch=compute_80,code=sm_80', + '-gencode', 'arch=compute_86,code=sm_86' +] + +setup( + name='serialization_cuda', + ext_modules=[ + CUDAExtension( + name='serialization_cuda', + sources=[ + 'cuda/bindings.cpp', + 'cuda/hilbert.cu', + 'cuda/morton.cu', + ], + extra_compile_args={ + 'cxx': ['-O3'], + 'nvcc': ['-O3']+all_cuda_archs + } + ) + ], + cmdclass={'build_ext': BuildExtension} +) \ No newline at end of file diff --git a/litept/serialization/z_order.py b/libs/serialization/z_order.py similarity index 100% rename from litept/serialization/z_order.py rename to libs/serialization/z_order.py diff --git a/litept/model.py b/litept/model.py index 267fb96..4b0a20f 100644 --- a/litept/model.py +++ b/litept/model.py @@ -16,7 +16,7 @@ from timm.layers import DropPath from libs.pointrope import PointROPE -from .serialization import encode +from libs.serialization import encode @torch.no_grad() diff --git a/models/utils/__init__.py b/models/utils/__init__.py index d4605f4..646f35e 100644 --- a/models/utils/__init__.py +++ b/models/utils/__init__.py @@ -5,4 +5,3 @@ batch2offset, off_diagonal, ) -from .serialization import encode, decode diff --git a/models/utils/serialization/__init__.py b/models/utils/serialization/__init__.py deleted file mode 100644 index 058c5e1..0000000 --- a/models/utils/serialization/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from .default import ( - encode, - decode, - z_order_encode, - z_order_decode, - hilbert_encode, - hilbert_decode, -) diff --git a/models/utils/serialization/default.py b/models/utils/serialization/default.py deleted file mode 100644 index 233051e..0000000 --- a/models/utils/serialization/default.py +++ /dev/null @@ -1,79 +0,0 @@ -import torch -from .z_order import xyz2key as z_order_encode_ -from .z_order import key2xyz as z_order_decode_ -from .hilbert import encode as hilbert_encode_ -from .hilbert import decode as hilbert_decode_ - -import numpy as np -from colorhash import ColorHash - - -def int_to_plotly_rgb(x): - """Convert 1D torch.Tensor of int into plotly-friendly RGB format. - This operation is deterministic on the int values. - """ - assert isinstance(x, torch.Tensor) - assert x.dim() == 1 - assert not x.is_floating_point() - x = x.cpu().long().numpy() - palette = np.array([ColorHash(i).rgb for i in range(x.max() + 1)]) - return palette[x] - -@torch.inference_mode() -def encode(grid_coord, batch=None, depth=16, order="z"): - assert order in {"z", "z-trans", "hilbert", "hilbert-trans"} - if order == "z": - code = z_order_encode(grid_coord, depth=depth) - # code = torch.argsort(code) - # rgb = int_to_plotly_rgb(code) - # print('a') - elif order == "z-trans": - code = z_order_encode(grid_coord[:, [1, 0, 2]], depth=depth) - elif order == "hilbert": - code = hilbert_encode(grid_coord, depth=depth) - # code = torch.argsort(code) - # rgb = int_to_plotly_rgb(code) - # print('a') - elif order == "hilbert-trans": - code = hilbert_encode(grid_coord[:, [1, 0, 2]], depth=depth) - else: - raise NotImplementedError - if batch is not None: - batch = batch.long() - code = batch << depth * 3 | code - return code - - -@torch.inference_mode() -def decode(code, depth=16, order="z"): - assert order in {"z", "hilbert"} - batch = code >> depth * 3 - code = code & ((1 << depth * 3) - 1) - if order == "z": - grid_coord = z_order_decode(code, depth=depth) - elif order == "hilbert": - grid_coord = hilbert_decode(code, depth=depth) - else: - raise NotImplementedError - return grid_coord, batch - - -def z_order_encode(grid_coord: torch.Tensor, depth: int = 16): - x, y, z = grid_coord[:, 0].long(), grid_coord[:, 1].long(), grid_coord[:, 2].long() - # we block the support to batch, maintain batched code in Point class - code = z_order_encode_(x, y, z, b=None, depth=depth) - return code - - -def z_order_decode(code: torch.Tensor, depth): - x, y, z = z_order_decode_(code, depth=depth) - grid_coord = torch.stack([x, y, z], dim=-1) # (N, 3) - return grid_coord - - -def hilbert_encode(grid_coord: torch.Tensor, depth: int = 16): - return hilbert_encode_(grid_coord, num_dims=3, num_bits=depth) - - -def hilbert_decode(code: torch.Tensor, depth: int = 16): - return hilbert_decode_(code, num_dims=3, num_bits=depth) diff --git a/models/utils/serialization/hilbert.py b/models/utils/serialization/hilbert.py deleted file mode 100644 index 7935657..0000000 --- a/models/utils/serialization/hilbert.py +++ /dev/null @@ -1,295 +0,0 @@ -import torch - - -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 binary2gray(binary, axis=-1): - """Convert an array of binary values into Gray codes. - - This uses the classic X ^ (X >> 1) trick to compute the Gray code. - - Parameters: - ----------- - binary: An ndarray of binary values. - - axis: The axis along which to compute the gray code. Default=-1. - - Returns: - -------- - Returns an ndarray of Gray codes. - """ - shifted = right_shift(binary, axis=axis) - - # Do the X ^ (X >> 1) trick. - gray = torch.logical_xor(binary, shifted) - - return gray - - -def gray2binary(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(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 decode(hilberts, num_dims, num_bits): - """Decode an array of Hilbert integers into locations in a hypercube. - - 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: - ------- - hilberts - An ndarray of Hilbert integers. Must be an integer dtype and - cannot have fewer bits than num_dims * num_bits. - - 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 unsigned integers with the same shape as hilberts - but with an additional dimension of size num_dims. - """ - - if num_dims * num_bits > 64: - raise ValueError( - """ - num_dims=%d and num_bits=%d for %d bits total, which can't be encoded - into a uint64. Are you sure you need that many points on your Hilbert - curve? - """ - % (num_dims, num_bits) - ) - - # Handle the case where we got handed a naked integer. - hilberts = torch.atleast_1d(hilberts) - - # Keep around the shape for later. - orig_shape = hilberts.shape - bitpack_mask = 2 ** torch.arange(0, 8).to(hilberts.device) - bitpack_mask_rev = bitpack_mask.flip(-1) - - # Treat each of the hilberts as a s equence of eight uint8. - # This treats all of the inputs as uint64 and makes things uniform. - hh_uint8 = ( - hilberts.ravel().type(torch.int64).view(torch.uint8).reshape((-1, 8)).flip(-1) - ) - - # Turn these lists of uints into lists of bits and then truncate to the size - # we actually need for using Skilling's procedure. - hh_bits = ( - hh_uint8.unsqueeze(-1) - .bitwise_and(bitpack_mask_rev) - .ne(0) - .byte() - .flatten(-2, -1)[:, -num_dims * num_bits :] - ) - - # Take the sequence of bits and Gray-code it. - gray = binary2gray(hh_bits) - - # There has got to be a better way to do this. - # I could index them differently, but the eventual packbits likes it this way. - gray = gray.reshape((-1, num_bits, num_dims)).swapaxes(1, 2) - - # Iterate backwards through the bits. - for bit in range(num_bits - 1, -1, -1): - # Iterate backwards through the dimensions. - for dim in range(num_dims - 1, -1, -1): - # 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]), - 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) - - # Pad back out to 64 bits. - extra_dims = 64 - num_bits - padded = torch.nn.functional.pad(gray, (extra_dims, 0), "constant", 0) - - # Now chop these up into blocks of 8. - locs_chopped = padded.flip(-1).reshape((-1, num_dims, 8, 8)) - - # Take those blocks and turn them unto uint8s. - # from IPython import embed; embed() - locs_uint8 = (locs_chopped * bitpack_mask).sum(3).squeeze().type(torch.uint8) - - # Finally, treat these as uint64s. - flat_locs = locs_uint8.view(torch.int64) - - # Return them in the expected shape. - return flat_locs.reshape((*orig_shape, num_dims)) diff --git a/models/utils/serialization/z_order.py b/models/utils/serialization/z_order.py deleted file mode 100644 index b0237c8..0000000 --- a/models/utils/serialization/z_order.py +++ /dev/null @@ -1,119 +0,0 @@ -import torch -from typing import Optional, Union - - -class KeyLUT: - def __init__(self): - r256 = torch.arange(256, dtype=torch.int64) - r512 = torch.arange(512, dtype=torch.int64) - zero = torch.zeros(256, dtype=torch.int64) - device = torch.device("cpu") - - self._encode = { - device: ( - self.xyz2key(r256, zero, zero, 8), - self.xyz2key(zero, r256, zero, 8), - self.xyz2key(zero, zero, r256, 8), - ) - } - self._decode = {device: self.key2xyz(r512, 9)} - - def encode_lut(self, device=torch.device("cpu")): - if device not in self._encode: - cpu = torch.device("cpu") - self._encode[device] = tuple(e.to(device) for e in self._encode[cpu]) - return self._encode[device] - - def decode_lut(self, device=torch.device("cpu")): - if device not in self._decode: - cpu = torch.device("cpu") - self._decode[device] = tuple(e.to(device) for e in self._decode[cpu]) - return self._decode[device] - - def xyz2key(self, x, y, z, depth): - key = torch.zeros_like(x) - for i in range(depth): - mask = 1 << i - key = ( - key - | ((x & mask) << (2 * i + 2)) - | ((y & mask) << (2 * i + 1)) - | ((z & mask) << (2 * i + 0)) - ) - return key - - def key2xyz(self, key, depth): - x = torch.zeros_like(key) - y = torch.zeros_like(key) - z = torch.zeros_like(key) - for i in range(depth): - x = x | ((key & (1 << (3 * i + 2))) >> (2 * i + 2)) - y = y | ((key & (1 << (3 * i + 1))) >> (2 * i + 1)) - z = z | ((key & (1 << (3 * i + 0))) >> (2 * i + 0)) - return x, y, z - - -_key_lut = KeyLUT() - - -def xyz2key( - x: torch.Tensor, - y: torch.Tensor, - z: torch.Tensor, - b: Optional[Union[torch.Tensor, int]] = None, - depth: int = 16, -): - r"""Encodes :attr:`x`, :attr:`y`, :attr:`z` coordinates to the shuffled keys - based on pre-computed look up tables. The speed of this function is much - faster than the method based on for-loop. - - Args: - x (torch.Tensor): The x coordinate. - y (torch.Tensor): The y coordinate. - z (torch.Tensor): The z coordinate. - b (torch.Tensor or int): The batch index of the coordinates, and should be - smaller than 32768. If :attr:`b` is :obj:`torch.Tensor`, the size of - :attr:`b` must be the same as :attr:`x`, :attr:`y`, and :attr:`z`. - depth (int): The depth of the shuffled key, and must be smaller than 17 (< 17). - """ - - EX, EY, EZ = _key_lut.encode_lut(x.device) - x, y, z = x.long(), y.long(), z.long() - - mask = 255 if depth > 8 else (1 << depth) - 1 - key = EX[x & mask] | EY[y & mask] | EZ[z & mask] - if depth > 8: - mask = (1 << (depth - 8)) - 1 - key16 = EX[(x >> 8) & mask] | EY[(y >> 8) & mask] | EZ[(z >> 8) & mask] - key = key16 << 24 | key - - if b is not None: - b = b.long() - key = b << 48 | key - - return key - - -def key2xyz(key: torch.Tensor, depth: int = 16): - r"""Decodes the shuffled key to :attr:`x`, :attr:`y`, :attr:`z` coordinates - and the batch index based on pre-computed look up tables. - - Args: - key (torch.Tensor): The shuffled key. - depth (int): The depth of the shuffled key, and must be smaller than 17 (< 17). - """ - - DX, DY, DZ = _key_lut.decode_lut(key.device) - x, y, z = torch.zeros_like(key), torch.zeros_like(key), torch.zeros_like(key) - - b = key >> 48 - key = key & ((1 << 48) - 1) - - n = (depth + 2) // 3 - for i in range(n): - k = key >> (i * 9) & 511 - x = x | (DX[k] << (i * 3)) - y = y | (DY[k] << (i * 3)) - z = z | (DZ[k] << (i * 3)) - - return x, y, z, b diff --git a/models/utils/structure.py b/models/utils/structure.py index afcc145..4772e0f 100644 --- a/models/utils/structure.py +++ b/models/utils/structure.py @@ -5,8 +5,7 @@ from addict import Dict from typing import List - -from models.utils.serialization import encode +from libs.serialization import encode from models.utils import ( offset2batch, batch2offset,