File size: 1,661 Bytes
178f950
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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

);