File size: 1,585 Bytes
f29eac5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
/*
 * Serialize a voxel grid
 *
 * Copyright (C) 2024, Jianfeng XIANG <[email protected]>
 * All rights reserved.
 *
 * Licensed under The MIT License [see LICENSE for details]
 *
 * Written by Jianfeng XIANG
 */

#pragma once
#include <torch/extension.h>


#define BLOCK_SIZE 256


/**
 * Z-order encode 3D points
 *
 * @param x [N] tensor containing the x coordinates
 * @param y [N] tensor containing the y coordinates
 * @param z [N] tensor containing the z coordinates
 *
 * @return [N] tensor containing the z-order encoded values
 */
torch::Tensor
z_order_encode(
    const torch::Tensor& x,
    const torch::Tensor& y,
    const torch::Tensor& z
);


/**
 * Z-order decode 3D points
 *
 * @param codes [N] tensor containing the z-order encoded values
 *
 * @return 3 tensors [N] containing the x, y, z coordinates
 */
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
z_order_decode(
    const torch::Tensor& codes
);


/**
 * Hilbert encode 3D points
 *
 * @param x [N] tensor containing the x coordinates
 * @param y [N] tensor containing the y coordinates
 * @param z [N] tensor containing the z coordinates
 *
 * @return [N] tensor containing the Hilbert encoded values
 */
torch::Tensor
hilbert_encode(
    const torch::Tensor& x,
    const torch::Tensor& y,
    const torch::Tensor& z
);


/**
 * Hilbert decode 3D points
 *
 * @param codes [N] tensor containing the Hilbert encoded values
 *
 * @return 3 tensors [N] containing the x, y, z coordinates
 */
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
hilbert_decode(
    const torch::Tensor& codes
);