souging/TRELLIS_TextTo3D

1#include <cuda.h>2#include <cuda_runtime.h>3#include <device_launch_parameters.h>4 5#include <cooperative_groups.h>6#include <cooperative_groups/memcpy_async.h>7namespace cg = cooperative_groups;8 9#include "hilbert.h"10 11 12// Expands a 10-bit integer into 30 bits by inserting 2 zeros after each bit.13static __device__ uint32_t expandBits(uint32_t v)14{15    v = (v * 0x00010001u) & 0xFF0000FFu;16    v = (v * 0x00000101u) & 0x0F00F00Fu;17    v = (v * 0x00000011u) & 0xC30C30C3u;18    v = (v * 0x00000005u) & 0x49249249u;19    return v;20}21 22 23// Removes 2 zeros after each bit in a 30-bit integer.24static __device__ uint32_t extractBits(uint32_t v)25{26    v = v & 0x49249249;27    v = (v ^ (v >>  2)) & 0x030C30C3u;28    v = (v ^ (v >>  4)) & 0x0300F00Fu;29    v = (v ^ (v >>  8)) & 0x030000FFu;30    v = (v ^ (v >> 16)) & 0x000003FFu;31    return v;32}33 34 35__global__ void hilbert_encode_cuda(36    size_t N,37    const uint32_t* x,38    const uint32_t* y,39    const uint32_t* z,40    uint32_t* codes41) {42    size_t thread_id = cg::this_grid().thread_rank();43    if (thread_id >= N) return;44 45    uint32_t point[3] = {x[thread_id], y[thread_id], z[thread_id]};46 47    uint32_t m = 1 << 9, q, p, t;48 49    // Inverse undo excess work50    q = m;51    while (q > 1) {52        p = q - 1;53        for (int i = 0; i < 3; i++) {54            if (point[i] & q) {55                point[0] ^= p;  // invert56            } else {57                t = (point[0] ^ point[i]) & p;58                point[0] ^= t;59                point[i] ^= t;60            }61        }62        q >>= 1;63    }64 65    // Gray encode66    for (int i = 1; i < 3; i++) {67        point[i] ^= point[i - 1];68    }69    t = 0;70    q = m;71    while (q > 1) {72        if (point[2] & q) {73            t ^= q - 1;74        }75        q >>= 1;76    }77    for (int i = 0; i < 3; i++) {78        point[i] ^= t;79    }80 81    // Convert to 3D Hilbert code82    uint32_t xx = expandBits(point[0]);83    uint32_t yy = expandBits(point[1]);84    uint32_t zz = expandBits(point[2]);85 86    codes[thread_id] = xx * 4 + yy * 2 + zz;87}88 89 90__global__ void hilbert_decode_cuda(91    size_t N,92    const uint32_t* codes,93    uint32_t* x,94    uint32_t* y,95    uint32_t* z96) {97    size_t thread_id = cg::this_grid().thread_rank();98    if (thread_id >= N) return;99 100    uint32_t point[3];101    point[0] = extractBits(codes[thread_id] >> 2);102    point[1] = extractBits(codes[thread_id] >> 1);103    point[2] = extractBits(codes[thread_id]);104 105    uint32_t m = 2 << 9, q, p, t;106 107    // Gray decode by H ^ (H/2)108    t = point[2] >> 1;109    for (int i = 2; i > 0; i--) {110        point[i] ^= point[i - 1];111    }112    point[0] ^= t;113 114    // Undo excess work115    q = 2;116    while (q != m) {117        p = q - 1;118        for (int i = 2; i >= 0; i--) {119            if (point[i] & q) {120                point[0] ^= p;121            } else {122                t = (point[0] ^ point[i]) & p;123                point[0] ^= t;124                point[i] ^= t;125            }126        }127        q <<= 1;128    }129 130    x[thread_id] = point[0];131    y[thread_id] = point[1];132    z[thread_id] = point[2];133}134