基于传统方式与BMI2指令集的3D Morton Z阶编解码C语言高效实现问询

Morton codes are a staple for spatial indexing, and getting their encoding/decoding right (and fast) can make a huge difference in performance for 3D grid-based applications. Let’s break down two approaches: the classic bit-interleaving method, and the optimized version using Intel’s BMI2 instruction set.

Traditional Bit-Interleaving Approach #

This method relies on shifting and masking to "spread" bits of each coordinate into their respective positions in the 64-bit Morton code, then reversing the process for decoding.

Encode Function #

First, we need a helper to spread a 32-bit coordinate into a 64-bit value with two zeros between each original bit. We then combine the spread x, y, z values by shifting them into their Morton positions (x in bits 0,3,6..., y in 1,4,7..., z in 2,5,8...).

#include <stdint.h>

// Helper: Spread 32-bit bits into 64-bit with two zeros between each bit
static uint64_t spread_bits(uint32_t w) {
    uint64_t x = w;
    x = (x | (x << 16)) & 0x0000FFFF0000FFFFULL;
    x = (x | (x << 8)) & 0x00FF00FF00FF00FFULL;
    x = (x | (x << 4)) & 0x0F0F0F0F0F0F0F0FULL;
    x = (x | (x << 2)) & 0x3333333333333333ULL;
    x = (x | (x << 1)) & 0x5555555555555555ULL;
    return x;
}

uint64_t morton_encode(uint32_t xindex, uint32_t yindex, uint32_t zindex) {
    return spread_bits(xindex) | (spread_bits(yindex) << 1) | (spread_bits(zindex) << 2);
}

Each step in spread_bits progressively inserts zeros between the original bits, doubling the number of gaps each time until we reach the 64-bit Morton-friendly format.

Decode Function #

To decode, we reverse the spreading process: extract each coordinate’s bits from the Morton code and compact them back into a 32-bit value.

// Helper: Compact 64-bit spread bits back into 32-bit
static uint32_t compact_bits(uint64_t x) {
    x &= 0x5555555555555555ULL;
    x = (x | (x >> 1)) & 0x3333333333333333ULL;
    x = (x | (x >> 2)) & 0x0F0F0F0F0F0F0F0FULL;
    x = (x | (x >> 4)) & 0x00FF00FF00FF00FFULL;
    x = (x | (x >> 8)) & 0x0000FFFF0000FFFFULL;
    x = (x | (x >> 16)) & 0x00000000FFFFFFFFULL;
    return (uint32_t)x;
}

void morton_decode(uint64_t morton_number, uint32_t *xindex, uint32_t *yindex, uint32_t *zindex) {
    *xindex = compact_bits(morton_number);
    *yindex = compact_bits(morton_number >> 1);
    *zindex = compact_bits(morton_number >> 2);
}

This function discards the zero bits inserted during encoding, packing the original coordinate bits back into a contiguous 32-bit value.

BMI2 Instruction-Based Optimizations #

Intel’s BMI2 instruction set (available on Haswell+ Intel CPUs and Zen+ AMD CPUs) introduces PDEP (Parallel Deposit) and PEXT (Parallel Extract) instructions that do bit spreading/compacting in a single CPU cycle. This is a massive speedup over the traditional method.

Encode with PDEP #

PDEP takes a source value and a mask, then deposits the source bits into the positions marked in the mask. For 3D Morton codes, we use masks that target the correct bit positions for each coordinate.

#include <stdint.h>
#include <immintrin.h>

uint64_t morton_encode(uint32_t xindex, uint32_t yindex, uint32_t zindex) {
    // Masks for 3D Morton positions: x (0,3,6...), y (1,4,7...), z (2,5,8...)
    const uint64_t x_mask = 0x1249249249249249ULL;
    const uint64_t y_mask = x_mask << 1;
    const uint64_t z_mask = x_mask << 2;

    return _pdep_u64(xindex, x_mask) | _pdep_u64(yindex, y_mask) | _pdep_u64(zindex, z_mask);
}

Each _pdep_u64 call directly places the coordinate’s bits into their Morton positions—no looping or repeated shifts needed.

Decode with PEXT #

PEXT does the reverse: it extracts bits from the source where the mask has bits set, then packs them into a contiguous value.

#include <stdint.h>
#include <immintrin.h>

void morton_decode(uint64_t morton_number, uint32_t *xindex, uint32_t *yindex, uint32_t *zindex) {
    const uint64_t x_mask = 0x1249249249249249ULL;
    const uint64_t y_mask = x_mask << 1;
    const uint64_t z_mask = x_mask << 2;

    *xindex = (uint32_t)_pext_u64(morton_number, x_mask);
    *yindex = (uint32_t)_pext_u64(morton_number, y_mask);
    *zindex = (uint32_t)_pext_u64(morton_number, z_mask);
}

This cuts decode time to just three single-cycle instructions—far faster than the traditional compacting steps.

Performance & Compatibility #

• Traditional Approach: Works on all CPUs, but takes ~15-20 cycles per encode/decode due to the series of shifts/masks.
• BMI2 Approach: Delivers a 3-6x speedup (~3-5 cycles total) on supported hardware. Note: You’ll need to enable BMI2 support in your compiler (e.g., -mbmi2 for GCC/Clang, /ARCH:AVX2 for MSVC).

For cross-platform code, add a fallback using preprocessor guards:

#ifdef __BMI2__
// BMI2 implementation here
#else
// Traditional implementation here
#endif

内容的提问来源于stack exchange，提问作者datapanda