explodethosebits 0.3.0__cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

include/etb/cuda/path_generator_kernel.cuh

@@ -0,0 +1,272 @@
+#ifndef ETB_PATH_GENERATOR_KERNEL_CUH
+#define ETB_PATH_GENERATOR_KERNEL_CUH
+
+#include "cuda_common.cuh"
+#include "gpu_memory.cuh"
+
+namespace etb {
+namespace cuda {
+
+/**
+ * Work item for path generation.
+ * Represents a partial path that needs to be explored.
+ */
+struct PathWorkItem {
+    uint32_t start_byte;        // Starting byte index for this work item
+    uint32_t current_depth;     // Current depth in the path
+    uint8_t prefix_bytes[16];   // Reconstructed bytes so far (max 16 for early stopping)
+    uint8_t prefix_length;      // Number of bytes in prefix
+    uint8_t bit_selections[16]; // Bit selections made so far
+    float current_score;        // Current heuristic score
+    
+    __host__ __device__ PathWorkItem()
+        : start_byte(0), current_depth(0), prefix_length(0), current_score(0.0f) {
+        for (int i = 0; i < 16; ++i) {
+            prefix_bytes[i] = 0;
+            bit_selections[i] = 0;
+        }
+    }
+};
+
+/**
+ * Configuration for path generator kernel.
+ */
+struct PathGeneratorConfig {
+    uint32_t input_length;          // Length of input data
+    uint32_t max_depth;             // Maximum path depth
+    uint32_t batch_size;            // Number of paths to generate per kernel launch
+    DeviceBitPruningConfig bit_pruning;
+    DeviceEarlyStoppingConfig early_stopping;
+    DeviceHeuristicWeights heuristic_weights;
+    DeviceScoringWeights scoring_weights;
+    
+    __host__ __device__ PathGeneratorConfig()
+        : input_length(0), max_depth(16), batch_size(65536) {}
+};
+
+/**
+ * Shared memory structure for path generation.
+ * Used for cooperative path exploration within a thread block.
+ */
+struct PathGeneratorSharedMem {
+    // Prefix state shared across warp
+    uint8_t shared_prefix[32];
+    uint32_t shared_prefix_length;
+    
+    // Work stealing queue (per-block)
+    uint32_t local_work_head;
+    uint32_t local_work_tail;
+    PathWorkItem local_work_items[32];  // Small local queue
+    
+    // Reduction scratch space
+    float warp_scores[32];
+    uint32_t warp_votes[32];
+};
+
+/**
+ * OPTIMIZED Path generator CUDA kernel.
+ * 
+ * Generates paths using work-stealing across thread blocks with
+ * warp-level cooperative path exploration.
+ * 
+ * Optimizations applied:
+ * - Lock-free work stealing with proper CAS semantics (no increment-then-check)
+ * - Warp-cooperative early stopping checks using ballot_sync
+ * - Coalesced bit extraction (all lanes read same byte, extract different bits)
+ * - Block-level work coordination to reduce global queue contention
+ * 
+ * Requirements: 9.3
+ * 
+ * @param input_data Input byte array
+ * @param config Kernel configuration
+ * @param work_queue Global work queue
+ * @param work_queue_head Head pointer for work queue
+ * @param work_queue_tail Tail pointer for work queue
+ * @param prefix_trie Prefix trie for pruning
+ * @param candidates Output candidate queue
+ * @param candidate_count Number of candidates found
+ * @param min_score Minimum score threshold
+ */
+__global__ void path_generator_kernel(
+    const uint8_t* input_data,
+    PathGeneratorConfig config,
+    PathWorkItem* work_queue,
+    uint32_t* work_queue_head,
+    uint32_t* work_queue_tail,
+    DevicePrefixTrieNode* prefix_trie,
+    DeviceCandidate* candidates,
+    uint32_t* candidate_count,
+    float* min_score
+);
+
+/**
+ * Initialize work queue with starting positions.
+ * 
+ * @param work_queue Work queue to initialize
+ * @param work_queue_tail Tail pointer
+ * @param input_length Length of input data
+ * @param bit_mask Bit mask for allowed positions
+ */
+__global__ void init_work_queue_kernel(
+    PathWorkItem* work_queue,
+    uint32_t* work_queue_tail,
+    uint32_t input_length,
+    uint8_t bit_mask
+);
+
+/**
+ * Host-side launcher for path generator kernel.
+ */
+class PathGeneratorKernel {
+public:
+    PathGeneratorKernel();
+    ~PathGeneratorKernel();
+    
+    /**
+     * Configure the kernel for a specific device.
+     * @param device_id CUDA device ID
+     */
+    void configure(int device_id);
+    
+    /**
+     * Launch the path generator kernel.
+     * @param mem GPU memory manager
+     * @param config Kernel configuration
+     * @param stream CUDA stream (nullptr for default)
+     */
+    void launch(GPUMemoryManager& mem, const PathGeneratorConfig& config, 
+                cudaStream_t stream = nullptr);
+    
+    /**
+     * Initialize work queue with starting positions.
+     * @param mem GPU memory manager
+     * @param input_length Length of input data
+     * @param bit_mask Bit mask for allowed positions
+     * @param stream CUDA stream
+     */
+    void init_work_queue(GPUMemoryManager& mem, uint32_t input_length,
+                         uint8_t bit_mask, cudaStream_t stream = nullptr);
+    
+    /**
+     * Get the kernel configuration.
+     */
+    const KernelConfig& get_config() const { return kernel_config_; }
+
+private:
+    KernelConfig kernel_config_;
+    bool configured_;
+};
+
+// Device functions for path generation
+
+/**
+ * Extract a bit from input data at the given coordinate.
+ */
+__device__ inline uint8_t extract_bit(const uint8_t* data, uint32_t byte_idx, uint8_t bit_pos) {
+    return (data[byte_idx] >> bit_pos) & 1;
+}
+
+/**
+ * Reconstruct a byte from 8 bit selections.
+ */
+__device__ inline uint8_t reconstruct_byte(const uint8_t* bits) {
+    uint8_t result = 0;
+    for (int i = 0; i < 8; ++i) {
+        result |= (bits[i] & 1) << i;
+    }
+    return result;
+}
+
+/**
+ * Check if a bit position is allowed by the mask.
+ */
+__device__ inline bool is_bit_allowed(uint8_t bit_pos, uint8_t mask) {
+    return (mask >> bit_pos) & 1;
+}
+
+/**
+ * Count allowed bits in mask.
+ */
+__device__ inline int count_allowed_bits(uint8_t mask) {
+    return __popc(static_cast<unsigned int>(mask));
+}
+
+/**
+ * Get the nth allowed bit position.
+ */
+__device__ inline uint8_t get_nth_allowed_bit(uint8_t mask, int n) {
+    int count = 0;
+    for (uint8_t i = 0; i < 8; ++i) {
+        if ((mask >> i) & 1) {
+            if (count == n) return i;
+            ++count;
+        }
+    }
+    return 0;
+}
+
+/**
+ * Warp-level vote for early termination.
+ * Returns true if majority of warp votes to terminate.
+ */
+__device__ inline bool warp_vote_terminate(bool should_terminate) {
+    unsigned int vote = __ballot_sync(0xFFFFFFFF, should_terminate);
+    return __popc(vote) > 16;  // More than half the warp
+}
+
+/**
+ * Warp-level reduction for finding best score.
+ */
+__device__ inline float warp_reduce_max(float val) {
+    for (int offset = 16; offset > 0; offset /= 2) {
+        float other = __shfl_down_sync(0xFFFFFFFF, val, offset);
+        val = fmaxf(val, other);
+    }
+    return __shfl_sync(0xFFFFFFFF, val, 0);
+}
+
+/**
+ * Atomic work stealing from global queue.
+ * Returns true if work was successfully stolen.
+ */
+__device__ inline bool steal_work(PathWorkItem* work_queue, 
+                                   uint32_t* head, uint32_t* tail,
+                                   uint32_t queue_capacity,
+                                   PathWorkItem& item) {
+    uint32_t old_head = atomicAdd(head, 1);
+    uint32_t current_tail = *tail;
+    
+    if (old_head < current_tail) {
+        item = work_queue[old_head % queue_capacity];
+        return true;
+    }
+    
+    // No work available, restore head
+    atomicSub(head, 1);
+    return false;
+}
+
+/**
+ * Push work item to global queue.
+ * Returns true if successfully pushed.
+ */
+__device__ inline bool push_work(PathWorkItem* work_queue,
+                                  uint32_t* tail,
+                                  uint32_t queue_capacity,
+                                  const PathWorkItem& item) {
+    uint32_t old_tail = atomicAdd(tail, 1);
+    
+    if (old_tail < queue_capacity) {
+        work_queue[old_tail] = item;
+        return true;
+    }
+    
+    // Queue full, restore tail
+    atomicSub(tail, 1);
+    return false;
+}
+
+} // namespace cuda
+} // namespace etb
+
+#endif // ETB_PATH_GENERATOR_KERNEL_CUH

include/etb/cuda/prefix_pruner_kernel.cuh

@@ -0,0 +1,370 @@
+#ifndef ETB_PREFIX_PRUNER_KERNEL_CUH
+#define ETB_PREFIX_PRUNER_KERNEL_CUH
+
+#include "cuda_common.cuh"
+#include "gpu_memory.cuh"
+
+namespace etb {
+namespace cuda {
+
+/**
+ * Shared memory structure for prefix pruning operations.
+ */
+struct PrefixPrunerSharedMem {
+    // Warp voting results
+    uint32_t warp_votes[8];  // 8 warps per block max
+    
+    // Prefix being evaluated
+    uint8_t current_prefix[16];
+    uint32_t prefix_length;
+    
+    // Trie navigation state
+    uint32_t current_node_idx;
+    DevicePrefixStatus current_status;
+    float current_score;
+};
+
+/**
+ * Prefix pruner CUDA kernel.
+ * 
+ * OPTIMIZED: Now takes pre-computed prefix offsets to avoid O(n²) offset calculation.
+ * Host should compute prefix_offsets using exclusive_scan(prefix_lengths).
+ * 
+ * Implements warp-level voting for termination decisions and
+ * atomic trie updates for prefix status.
+ * 
+ * Requirements: 9.8
+ * 
+ * @param prefix_trie Prefix trie nodes
+ * @param trie_size Number of nodes in trie
+ * @param prefixes Array of prefixes to evaluate
+ * @param prefix_lengths Array of prefix lengths
+ * @param prefix_offsets Pre-computed exclusive prefix sum of lengths (NEW)
+ * @param scores Array of heuristic scores for each prefix
+ * @param num_prefixes Number of prefixes to evaluate
+ * @param prune_threshold Score threshold for pruning
+ * @param prune_results Output array indicating if each prefix was pruned
+ */
+__global__ void prefix_pruner_kernel(
+    DevicePrefixTrieNode* prefix_trie,
+    uint32_t trie_size,
+    const uint8_t* prefixes,
+    const uint32_t* prefix_lengths,
+    const uint32_t* prefix_offsets,
+    const float* scores,
+    uint32_t num_prefixes,
+    float prune_threshold,
+    bool* prune_results
+);
+
+/**
+ * Trie lookup kernel.
+ * OPTIMIZED: Uses pre-computed prefix offsets.
+ * 
+ * @param prefix_trie Prefix trie nodes
+ * @param trie_size Number of nodes in trie
+ * @param prefixes Array of prefixes to look up
+ * @param prefix_lengths Array of prefix lengths
+ * @param prefix_offsets Pre-computed exclusive prefix sum of lengths (NEW)
+ * @param num_prefixes Number of prefixes
+ * @param statuses Output array of prefix statuses
+ * @param scores Output array of best scores
+ */
+__global__ void trie_lookup_kernel(
+    const DevicePrefixTrieNode* prefix_trie,
+    uint32_t trie_size,
+    const uint8_t* prefixes,
+    const uint32_t* prefix_lengths,
+    const uint32_t* prefix_offsets,
+    uint32_t num_prefixes,
+    DevicePrefixStatus* statuses,
+    float* scores
+);
+
+/**
+ * Trie insert/update kernel.
+ * OPTIMIZED: Uses pre-computed prefix offsets and warp-cooperative allocation.
+ * 
+ * @param prefix_trie Prefix trie nodes
+ * @param trie_size Current number of nodes
+ * @param max_trie_size Maximum trie capacity
+ * @param prefixes Array of prefixes to insert
+ * @param prefix_lengths Array of prefix lengths
+ * @param prefix_offsets Pre-computed exclusive prefix sum of lengths (NEW)
+ * @param statuses Array of statuses to set
+ * @param scores Array of scores to set
+ * @param num_prefixes Number of prefixes
+ * @param new_trie_size Output: new trie size after insertions
+ */
+__global__ void trie_insert_kernel(
+    DevicePrefixTrieNode* prefix_trie,
+    uint32_t trie_size,
+    uint32_t max_trie_size,
+    const uint8_t* prefixes,
+    const uint32_t* prefix_lengths,
+    const uint32_t* prefix_offsets,
+    const DevicePrefixStatus* statuses,
+    const float* scores,
+    uint32_t num_prefixes,
+    uint32_t* new_trie_size
+);
+
+/**
+ * Batch prune check kernel.
+ * OPTIMIZED: Uses pre-computed prefix offsets.
+ * 
+ * @param prefix_trie Prefix trie nodes
+ * @param trie_size Number of nodes
+ * @param prefixes Array of prefixes to check
+ * @param prefix_lengths Array of prefix lengths
+ * @param prefix_offsets Pre-computed exclusive prefix sum of lengths (NEW)
+ * @param num_prefixes Number of prefixes
+ * @param should_skip Output array indicating if prefix should be skipped
+ */
+__global__ void batch_prune_check_kernel(
+    const DevicePrefixTrieNode* prefix_trie,
+    uint32_t trie_size,
+    const uint8_t* prefixes,
+    const uint32_t* prefix_lengths,
+    const uint32_t* prefix_offsets,
+    uint32_t num_prefixes,
+    bool* should_skip
+);
+
+/**
+ * Host-side launcher for prefix pruner kernels.
+ */
+class PrefixPrunerKernel {
+public:
+    PrefixPrunerKernel();
+    ~PrefixPrunerKernel();
+    
+    /**
+     * Configure the kernel for a specific device.
+     * @param device_id CUDA device ID
+     */
+    void configure(int device_id);
+    
+    /**
+     * Evaluate and prune prefixes.
+     * @param trie Device prefix trie
+     * @param trie_size Number of trie nodes
+     * @param prefixes Device array of prefixes (flattened)
+     * @param prefix_lengths Device array of prefix lengths
+     * @param scores Device array of heuristic scores
+     * @param num_prefixes Number of prefixes
+     * @param prune_threshold Score threshold for pruning
+     * @param prune_results Device array for results
+     * @param stream CUDA stream
+     */
+    void evaluate_and_prune(DevicePrefixTrieNode* trie, uint32_t trie_size,
+                            const uint8_t* prefixes, const uint32_t* prefix_lengths,
+                            const float* scores, uint32_t num_prefixes,
+                            float prune_threshold, bool* prune_results,
+                            cudaStream_t stream = nullptr);
+    
+    /**
+     * Look up prefix statuses.
+     * @param trie Device prefix trie
+     * @param trie_size Number of trie nodes
+     * @param prefixes Device array of prefixes
+     * @param prefix_lengths Device array of prefix lengths
+     * @param num_prefixes Number of prefixes
+     * @param statuses Device array for output statuses
+     * @param scores Device array for output scores
+     * @param stream CUDA stream
+     */
+    void lookup(const DevicePrefixTrieNode* trie, uint32_t trie_size,
+                const uint8_t* prefixes, const uint32_t* prefix_lengths,
+                uint32_t num_prefixes, DevicePrefixStatus* statuses,
+                float* scores, cudaStream_t stream = nullptr);
+    
+    /**
+     * Check if prefixes should be skipped due to pruned ancestors.
+     * @param trie Device prefix trie
+     * @param trie_size Number of trie nodes
+     * @param prefixes Device array of prefixes
+     * @param prefix_lengths Device array of prefix lengths
+     * @param num_prefixes Number of prefixes
+     * @param should_skip Device array for output flags
+     * @param stream CUDA stream
+     */
+    void check_pruned(const DevicePrefixTrieNode* trie, uint32_t trie_size,
+                      const uint8_t* prefixes, const uint32_t* prefix_lengths,
+                      uint32_t num_prefixes, bool* should_skip,
+                      cudaStream_t stream = nullptr);
+    
+    /**
+     * Get the kernel configuration.
+     */
+    const KernelConfig& get_config() const { return kernel_config_; }
+
+private:
+    KernelConfig kernel_config_;
+    bool configured_;
+};
+
+// ============================================================================
+// Device Functions
+// ============================================================================
+
+/**
+ * Navigate trie to find node for a prefix.
+ * Returns node index or UINT32_MAX if not found.
+ */
+__device__ inline uint32_t find_prefix_node(
+    const DevicePrefixTrieNode* trie,
+    uint32_t trie_size,
+    const uint8_t* prefix,
+    uint32_t prefix_length
+) {
+    if (trie_size == 0 || prefix_length == 0) {
+        return 0;  // Root node
+    }
+    
+    uint32_t current = 0;  // Start at root
+    
+    for (uint32_t i = 0; i < prefix_length; ++i) {
+        uint8_t byte_val = prefix[i];
+        uint32_t children_offset = trie[current].children_offset;
+        
+        if (children_offset == 0) {
+            return UINT32_MAX;  // No children, prefix not found
+        }
+        
+        // Look for child with matching byte value
+        bool found = false;
+        for (uint32_t c = 0; c < 256 && children_offset + c < trie_size; ++c) {
+            uint32_t child_idx = children_offset + c;
+            if (child_idx < trie_size && trie[child_idx].reconstructed_byte == byte_val) {
+                current = child_idx;
+                found = true;
+                break;
+            }
+        }
+        
+        if (!found) {
+            return UINT32_MAX;
+        }
+    }
+    
+    return current;
+}
+
+/**
+ * Check if any ancestor of a prefix is pruned.
+ */
+__device__ inline bool is_ancestor_pruned(
+    const DevicePrefixTrieNode* trie,
+    uint32_t trie_size,
+    const uint8_t* prefix,
+    uint32_t prefix_length
+) {
+    if (trie_size == 0) return false;
+    
+    uint32_t current = 0;
+    
+    for (uint32_t i = 0; i < prefix_length; ++i) {
+        // Check current node status
+        if (trie[current].status == DevicePrefixStatus::PRUNED) {
+            return true;
+        }
+        
+        uint8_t byte_val = prefix[i];
+        uint32_t children_offset = trie[current].children_offset;
+        
+        if (children_offset == 0) {
+            return false;  // No more nodes to check
+        }
+        
+        // Find child
+        bool found = false;
+        for (uint32_t c = 0; c < 256 && children_offset + c < trie_size; ++c) {
+            uint32_t child_idx = children_offset + c;
+            if (child_idx < trie_size && trie[child_idx].reconstructed_byte == byte_val) {
+                current = child_idx;
+                found = true;
+                break;
+            }
+        }
+        
+        if (!found) {
+            return false;
+        }
+    }
+    
+    // Check final node
+    return trie[current].status == DevicePrefixStatus::PRUNED;
+}
+
+/**
+ * Atomically update node status.
+ */
+__device__ inline void atomic_update_status(
+    DevicePrefixTrieNode* node,
+    DevicePrefixStatus new_status
+) {
+    // Use atomicCAS on the status byte
+    uint8_t* status_ptr = reinterpret_cast<uint8_t*>(&node->status);
+    uint8_t old_val = *status_ptr;
+    uint8_t new_val = static_cast<uint8_t>(new_status);
+    
+    // Only update if transitioning to a "more final" state
+    // UNKNOWN -> VALID or PRUNED is allowed
+    // VALID -> PRUNED is allowed
+    // PRUNED is final
+    if (old_val == static_cast<uint8_t>(DevicePrefixStatus::PRUNED)) {
+        return;  // Already pruned, don't change
+    }
+    
+    atomicCAS(reinterpret_cast<unsigned int*>(status_ptr), 
+              static_cast<unsigned int>(old_val),
+              static_cast<unsigned int>(new_val));
+}
+
+/**
+ * Atomically update best score (only if new score is higher).
+ */
+__device__ inline void atomic_update_score(
+    DevicePrefixTrieNode* node,
+    float new_score
+) {
+    // Use atomicMax on the score
+    // Since atomicMax doesn't work directly on floats, we use a CAS loop
+    float* score_ptr = &node->best_score;
+    float old_score = *score_ptr;
+    
+    while (new_score > old_score) {
+        float assumed = old_score;
+        old_score = __int_as_float(atomicCAS(
+            reinterpret_cast<int*>(score_ptr),
+            __float_as_int(assumed),
+            __float_as_int(new_score)
+        ));
+        
+        if (old_score == assumed) {
+            break;  // Successfully updated
+        }
+    }
+}
+
+/**
+ * Warp-level vote for pruning decision.
+ * Returns true if majority of warp votes to prune.
+ */
+__device__ inline bool warp_vote_prune(bool should_prune) {
+    unsigned int vote = __ballot_sync(0xFFFFFFFF, should_prune);
+    return __popc(vote) > 16;  // More than half
+}
+
+/**
+ * Increment visit count atomically.
+ */
+__device__ inline void atomic_increment_visit(DevicePrefixTrieNode* node) {
+    atomicAdd(&node->visit_count, 1);
+}
+
+} // namespace cuda
+} // namespace etb
+
+#endif // ETB_PREFIX_PRUNER_KERNEL_CUH