souging/TRELLIS_TextTo3D

1#include <torch/extension.h>2#include "api.h"3#include "z_order.h"4#include "hilbert.h"5 6 7torch::Tensor8z_order_encode(9    const torch::Tensor& x,10    const torch::Tensor& y,11    const torch::Tensor& z12) {13    // Allocate output tensor14    torch::Tensor codes = torch::empty_like(x);15 16    // Call CUDA kernel17    z_order_encode_cuda<<<(x.size(0) + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE>>>(18        x.size(0),19        reinterpret_cast<uint32_t*>(x.contiguous().data_ptr<int>()),20        reinterpret_cast<uint32_t*>(y.contiguous().data_ptr<int>()),21        reinterpret_cast<uint32_t*>(z.contiguous().data_ptr<int>()),22        reinterpret_cast<uint32_t*>(codes.data_ptr<int>())23    );24 25    return codes;26}27 28 29std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>30z_order_decode(31    const torch::Tensor& codes32) {33    // Allocate output tensors34    torch::Tensor x = torch::empty_like(codes);35    torch::Tensor y = torch::empty_like(codes);36    torch::Tensor z = torch::empty_like(codes);37 38    // Call CUDA kernel39    z_order_decode_cuda<<<(codes.size(0) + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE>>>(40        codes.size(0),41        reinterpret_cast<uint32_t*>(codes.contiguous().data_ptr<int>()),42        reinterpret_cast<uint32_t*>(x.data_ptr<int>()),43        reinterpret_cast<uint32_t*>(y.data_ptr<int>()),44        reinterpret_cast<uint32_t*>(z.data_ptr<int>())45    );46 47    return std::make_tuple(x, y, z);48}49 50 51torch::Tensor52hilbert_encode(53    const torch::Tensor& x,54    const torch::Tensor& y,55    const torch::Tensor& z56) {57    // Allocate output tensor58    torch::Tensor codes = torch::empty_like(x);59 60    // Call CUDA kernel61    hilbert_encode_cuda<<<(x.size(0) + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE>>>(62        x.size(0),63        reinterpret_cast<uint32_t*>(x.contiguous().data_ptr<int>()),64        reinterpret_cast<uint32_t*>(y.contiguous().data_ptr<int>()),65        reinterpret_cast<uint32_t*>(z.contiguous().data_ptr<int>()),66        reinterpret_cast<uint32_t*>(codes.data_ptr<int>())67    );68 69    return codes;70}71 72 73std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>74hilbert_decode(75    const torch::Tensor& codes76) {77    // Allocate output tensors78    torch::Tensor x = torch::empty_like(codes);79    torch::Tensor y = torch::empty_like(codes);80    torch::Tensor z = torch::empty_like(codes);81 82    // Call CUDA kernel83    hilbert_decode_cuda<<<(codes.size(0) + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE>>>(84        codes.size(0),85        reinterpret_cast<uint32_t*>(codes.contiguous().data_ptr<int>()),86        reinterpret_cast<uint32_t*>(x.data_ptr<int>()),87        reinterpret_cast<uint32_t*>(y.data_ptr<int>()),88        reinterpret_cast<uint32_t*>(z.data_ptr<int>())89    );90 91    return std::make_tuple(x, y, z);92}93