Feat/gpu wavelet codecs 64bit

HIP_API.md

@@ -6,11 +6,13 @@
 > format. Compression ratios differ from the CPU reference due to different
 > block tiling strategies.
 
-GPU-accelerated lossy compression for 3D floating-point volumes on AMD Instinct
-GPUs (MI200, MI300). Targets seismic imaging workloads where wavefield snapshots
-must be stored and retrieved at GPU memory bandwidth.
+GPU-accelerated lossy compression for 2D and 3D floating-point volumes on AMD
+Instinct GPUs (MI200, MI300, MI355). Targets seismic imaging workloads where
+wavefield snapshots must be stored and retrieved at GPU memory bandwidth.
 
-Single fused kernel: wavelet transform (DS 7/9) → quantization → RLE encoding.
+Single fused kernel: wavelet transform (DS 7/9) → quantization → significance/RLE
+coding. The coder is selectable per plan (see Kernel Variants); the default
+resolves to octree for 3D and quadtree for 2D.
 Error norms match the CPU reference (CvxCompress) to floating-point rounding.
 
 ### Performance (MI300X vs 128-core EPYC 9554, AVX, best thread count)
@@ -27,14 +29,14 @@ time rather than end-to-end latency.
 
 ## Requirements
 
-- ROCm 7.x (`module load rocm/7.2.0`)
-- AMD GPU: gfx90a (MI200) or gfx942 (MI300X)
+- ROCm 7.x (`module load rocm/7.2.1`)
+- AMD GPU: gfx90a (MI200), gfx942 (MI300X), or gfx950 (MI355X)
 - C++17, `hipcc`, `rocprim`
 
 ## Building
 
 ```bash
-module load rocm/7.2.0
+module load rocm/7.2.1
 
 # Build the CPU reference library (needed by tests)
 make libcvxcompress.so
@@ -71,9 +73,9 @@ All functions are declared in [`hip/hipCompress.h`](hip/hipCompress.h).
 
 | Function | Description |
 |----------|-------------|
-| `hipCompress` | Wavelet + quantize + RLE encode → self-contained compressed stream (async) |
+| `hipCompress` | Wavelet + quantize + encode (per-plan codec) → self-contained compressed stream (async) |
 | `hipCompressSynchronize` | Block until compress completes, retrieve compressed length and CR |
-| `hipDecompress` | RLE decode + inverse wavelet → wavelet buffer (single kernel, async) |
+| `hipDecompress` | Decode (per-plan codec) + inverse wavelet → wavelet buffer (single kernel, async) |
 
 ### Utilities
 
@@ -187,17 +189,81 @@ The `scale` argument to `hipCompress` controls the quality/compression tradeoff:
 
 ### Kernel Variants
 
-| Variant | Enum | Description |
-|---------|------|-------------|
-| Z-line  | `HIP_COMPRESS_KERNEL_ZLINE` | Parallel z-line RLE with per-block metadata (default) |
-| Seg-RLE | `HIP_COMPRESS_KERNEL_SEGRLE` | Segment-aligned RLE, no metadata overhead (unoptimized) |
+| Variant | Enum | Dims | Description |
+|---------|------|------|-------------|
+| Auto      | `HIP_COMPRESS_KERNEL_AUTO`     | 2D/3D | **Default.** Dimensionality-selected: resolves to quadtree for 2D (`nz == 1`) and octree for 3D at plan creation, falling back to z-line for configurations the structured coders cannot handle |
+| Z-line    | `HIP_COMPRESS_KERNEL_ZLINE`    | 2D/3D | Parallel z-line RLE with per-block metadata. Highest encode throughput; the throughput hedge for encode-bound paths (e.g. per-timestep RTM checkpoint spilling) |
+| Seg-RLE   | `HIP_COMPRESS_KERNEL_SEGRLE`   | 3D    | Segment-aligned RLE, no metadata overhead (unoptimized) |
+| Octree    | `HIP_COMPRESS_KERNEL_OCTREE`   | 3D    | Octree significance coder with per-block PFOR value coding. Best compression ratio and fastest decode; recommended for storage/archival |
+| Quadtree  | `HIP_COMPRESS_KERNEL_QUADTREE` | 2D    | Quadtree significance coder with per-block PFOR value coding -- the 2D counterpart of octree |
+| Two-level | `HIP_COMPRESS_KERNEL_TWOLEVEL` | 3D    | Two-level occupancy + per-line width coder. Higher encode throughput than octree at a lower ratio; a hedge for compute/bandwidth-bound, rewrite-heavy paths |
 
-Select at plan creation: `hipCompressCreatePlan(&plan, nx, ny, nz, aux, HIP_COMPRESS_KERNEL_SEGRLE)`.
+The codec is selected **at runtime, per plan** via the last argument of
+`hipCompressCreatePlan` — it is an ordinary function parameter, so switching
+between codecs requires **no recompilation** of the library or the application:
+
+```cpp
+// default (kernel arg omitted) → auto: octree for 3D, quadtree for 2D
+hipCompressCreatePlan(&plan_def, nx, ny, nz, aux);
+
+// storage-oriented volume → octree (best ratio) — explicit form of the 3D default
+hipCompressCreatePlan(&plan_oct, nx, ny, nz, aux, HIP_COMPRESS_KERNEL_OCTREE);
+
+// encode-bound / high-rate path → z-line (fastest encode)
+hipCompressCreatePlan(&plan_zl,  nx, ny, nz, aux, HIP_COMPRESS_KERNEL_ZLINE);
+```
+
+The codec is bound to the plan (internal buffer sizes and stream header layout
+differ per codec), so the switching granularity is "which plan you create"; an
+existing plan's codec cannot be changed in place. Octree and two-level are 3D
+only, quadtree is 2D only — requesting one for the wrong dimensionality fails
+plan creation with `HIP_COMPRESS_ERROR_INVALID_DIMENSIONS`. `AUTO` avoids this
+by resolving to the dimensionality-appropriate coder at plan creation.
+
+**Choosing a codec.** The `AUTO` default maximizes compression ratio and decode
+throughput and is the right choice for storage/archival and read-heavy paths.
+For **encode-bound** paths that compress on a hot loop — e.g. per-timestep RTM
+checkpoint spilling — prefer `ZLINE`, which has the highest encode throughput.
+The octree/quadtree encode cost over z-line is small at production grid sizes
+(~1–3% at 512³) but grows at small grids where the per-block histogram, scans,
+and PFOR bookkeeping are not amortized.
+
+**Memory footprint.** Octree/quadtree allocate a larger per-block scratch stride
+(`WOCT_CODE_SLOT_BYTES` ~140 KB/block plus the bitmap scratch) than z-line;
+these are global HBM buffers (not LDS). Negligible on MI300X/MI355X but worth
+noting for very large volumes.
+
+#### Two-level: architecture-selected encoder
+
+The two-level codec has two encoders that emit a **byte-identical** stream:
+
+- an LDS-staged **opt** encoder (~131 KB LDS) that requires CDNA4 (gfx950,
+  e.g. MI355x), and
+- a **portable** encoder used on every other architecture (gfx942/MI300x,
+  gfx90a, …).
+
+The encoder is chosen automatically at plan creation from the device
+architecture (`gcnArchName`); no user action is required. Because the streams
+are identical, a volume encoded on one GPU decodes correctly on any other, and
+the single-arch-independent decoder runs everywhere.
+
+For this automatic selection to reach the fast path, the library `.so` must
+have been built with the target architecture(s) included — build once as a fat
+binary covering your deployment GPUs and no per-user recompile is ever needed:
+
+```bash
+make libhipcvxcompress.so HIP_ARCH="gfx90a gfx942 gfx950"
+```
+
+`HIP_ARCH` accepts a space-separated list; each architecture is expanded to a
+`--offload-arch=` flag. The opt encoder is compiled only into the gfx950 device
+image (a no-op stub elsewhere), so the fat binary carries both, and the HIP
+loader plus the runtime dispatch pick the correct one for the GPU in use.
 
 ### Two-Stream Model
 
-- **`user_stream`**: passed to each API call. Wavelet transform and RLE encoding
-  run here. The stream is free immediately after `hipCompress` returns.
+- **`user_stream`**: passed to each API call. The wavelet transform and value
+  encoding run here. The stream is free immediately after `hipCompress` returns.
 - **`aux_stream`**: owned by the user, passed at plan creation. Compaction, header
   writing, and D2H readback run here. Shared across plans.
 
@@ -237,15 +303,17 @@ hip/
   hipWaveletRLEInverse.h         Fused inverse RLE + wavelet kernels
   hipRLEDecode.h                 Z-line RLE decoder
   hipSegmentedRLE.h              Segment-aligned RLE encode/decode
+  hipWaveletBitmap.h             Bitmap significance split + two-level coder/decoder
+  hipWaveletOctree.h             Octree significance coder + shared inverse stage
   ds79.h                         DS 7/9 wavelet filter coefficients and transforms
   ds79_reg32.inc                 Unrolled forward wavelet (32-point)
   us79_reg32.inc                 Unrolled inverse wavelet (32-point)
 tests/
-  test_compress_api_hip.cpp      API test suite (37 tests + benchmarks)
+  test_compress_api_hip.cpp      API test suite (39 tests + benchmarks)
   example_async_pipeline.cpp     Async overlap example
 ```
 
 ## License
 
-Copyright (C) 2025 Advanced Micro Devices, Inc. Licensed under the
+Copyright (C) 2026 Advanced Micro Devices, Inc. Licensed under the
 [MIT License](https://opensource.org/licenses/MIT).

README.md

@@ -18,12 +18,13 @@ https://library.seg.org/doi/pdf/10.1190/1.1826518
 > - Optimizations are ongoing
 > - Not tested in an integrated production setting
 > - No backward compatibility guarantee for the compressed bitstream format
-> - AMD GPUs only (ROCm 7.x, gfx90a/gfx942)
+> - AMD GPUs only (ROCm 7.x, gfx90a/gfx942/gfx950)
 > - Compression ratios differ from the CPU reference due to different block tiling strategies
 
 See [HIP_API.md](HIP_API.md) for the AMD GPU port of this library targeting
-MI200/MI300X via HIP. It implements a fused wavelet + quantization + RLE pipeline
-in single GPU kernels, achieving 14–27x speedup over the fully-parallelized CPU
+MI200/MI300X/MI355X via HIP. It implements a fused wavelet + quantization +
+significance/RLE coding pipeline in single GPU kernels (default: octree for 3D,
+quadtree for 2D), achieving 14–27x speedup over the fully-parallelized CPU
 reference (128-core EPYC 9554, AVX, best thread count) with matching error norms
 (to floating-point rounding). Includes API reference, usage examples, and async
 pipeline integration.

hip/ds79.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2025 Advanced Micro Devices, Inc.
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
 // Use of this source code is governed by an MIT-style license that can be
 // found in the LICENSE file or at https://opensource.org/licenses/MIT.
 

hip/ds79_f4_reg32.inc

@@ -1,4 +1,4 @@
-// Copyright (C) 2025 Advanced Micro Devices, Inc.
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
 // Use of this source code is governed by an MIT-style license that can be
 // found in the LICENSE file or at https://opensource.org/licenses/MIT.
 

hip/ds79_reg32.inc

@@ -1,4 +1,4 @@
-// Copyright (C) 2025 Advanced Micro Devices, Inc.
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
 // Use of this source code is governed by an MIT-style license that can be
 // found in the LICENSE file or at https://opensource.org/licenses/MIT.
 

hip/hipBlockCopy.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2025 Advanced Micro Devices, Inc.
+// Copyright (C) 2026 Advanced Micro Devices, Inc.
 // Use of this source code is governed by an MIT-style license that can be
 // found in the LICENSE file or at https://opensource.org/licenses/MIT.
 
@@ -7,6 +7,7 @@
 
 #include <hip/hip_runtime.h>
 #include <rocprim/block/block_reduce.hpp>
+#include "hipPlaneIO.h"
 
 using bcopy_float4_vec = __attribute__((__vector_size__(4 * sizeof(float)))) float;
 using bcopy_int4_vec   = __attribute__((__vector_size__(4 * sizeof(int)))) int;
@@ -47,23 +48,32 @@ __global__ void copyToWaveletKernelOpt(
 
     bool y_in_range = (gy < ey);
 
-    uint32_t src_row_byte = ((y0 + gy) * ldimx + x0 + gx) * (uint32_t)sizeof(float);
+    size_t src_row_byte = ((size_t)(y0 + gy) * ldimx + x0 + gx) * sizeof(float);
+    bool src_x_all_in = (gx + 3 < ex);
 
     // ---- Phase 1: Load ZPB planes (or zero-fill beyond extraction) ----
+    // Global loads have no OOB clamp (unlike the former buffer path), so the
+    // ragged x edge is read lane-by-lane, guarded by the extraction window.
     bcopy_float4_vec regs[BCOPY_ZPB];
     constexpr bcopy_float4_vec zero_vec = {0.0f, 0.0f, 0.0f, 0.0f};
     #pragma unroll
     for (int dz = 0; dz < BCOPY_ZPB; ++dz) {
         int iz = z_start + dz;
-        if (y_in_range && iz < ez) {
-            auto plane_rsrc = __builtin_amdgcn_make_buffer_rsrc(
-                const_cast<float*>(d_src + (long)(z0 + iz) * ldimxy),
-                0, -1, 0x00027000);
+        if (!(y_in_range && iz < ez) || gx >= ex) {
+            regs[dz] = zero_vec;
+            continue;
+        }
+        const float* plane = d_src + (size_t)(z0 + iz) * ldimxy;
+        if (src_x_all_in) {
             regs[dz] = __builtin_bit_cast(bcopy_float4_vec,
-                __builtin_amdgcn_raw_buffer_load_b128(
-                    plane_rsrc, src_row_byte, 0, SLC));
+                hipPlaneLoadF4u(plane, src_row_byte));
         } else {
-            regs[dz] = zero_vec;
+            bcopy_float4_vec v = zero_vec;
+            if (gx + 0 < ex) v[0] = hipPlaneLoadScalarNT(plane, src_row_byte + 0);
+            if (gx + 1 < ex) v[1] = hipPlaneLoadScalarNT(plane, src_row_byte + 4);
+            if (gx + 2 < ex) v[2] = hipPlaneLoadScalarNT(plane, src_row_byte + 8);
+            if (gx + 3 < ex) v[3] = hipPlaneLoadScalarNT(plane, src_row_byte + 12);
+            regs[dz] = v;
         }
     }
 
@@ -84,18 +94,14 @@ __global__ void copyToWaveletKernelOpt(
 
     // ---- Phase 3: Store ZPB planes (skip planes beyond wnz) ----
     if constexpr (DO_COPY) {
-        uint32_t dst_byte = (gx + gy * wnx) * (uint32_t)sizeof(float);
+        // Wavelet dst is 32-padded (wnx,wny multiples of 32) → always full tiles.
+        size_t dst_byte = ((size_t)gy * wnx + gx) * sizeof(float);
 
         #pragma unroll
         for (int dz = 0; dz < BCOPY_ZPB; ++dz) {
-            if (z_start + dz < wnz) {
-                auto plane_rsrc = __builtin_amdgcn_make_buffer_rsrc(
-                    d_dst + (long)(z_start + dz) * wnx * wny,
-                    0, -1, 0x00027000);
-                auto vi = __builtin_bit_cast(bcopy_int4_vec, regs[dz]);
-                __builtin_amdgcn_raw_buffer_store_b128(
-                    vi, plane_rsrc, dst_byte, 0, SLC);
-            }
+            if (z_start + dz < wnz)
+                hipPlaneStoreNT<bcopy_float4_vec>(
+                    d_dst + (size_t)(z_start + dz) * wnx * wny, dst_byte, regs[dz]);
         }
     }
 
@@ -173,59 +179,46 @@ __global__ void copyFromWaveletKernelOpt(
     if (gy >= ey) return;
     if (gx >= ex) return;
 
-    uint32_t src_byte = (gx + gy * wnx) * (uint32_t)sizeof(float);
-    long wav_plane = (long)wnx * wny;
+    // Wavelet src is 32-padded and the float4 never crosses a 32-tile boundary
+    // (xg*4+3 <= 31 < wnx), so the vector load is always fully in bounds.
+    size_t src_byte = ((size_t)gy * wnx + gx) * sizeof(float);
+    size_t wav_plane = (size_t)wnx * wny;
 
     constexpr bcopy_float4_vec zero_vec_from = {0.0f, 0.0f, 0.0f, 0.0f};
     bcopy_float4_vec regs[BCOPY_ZPB];
     #pragma unroll
     for (int dz = 0; dz < BCOPY_ZPB; ++dz) {
-        if (z_start + dz < wnz) {
-            auto plane_rsrc = __builtin_amdgcn_make_buffer_rsrc(
-                const_cast<float*>(d_src + (long)(z_start + dz) * wav_plane),
-                0, -1, 0x00027000);
-            regs[dz] = __builtin_bit_cast(bcopy_float4_vec,
-                __builtin_amdgcn_raw_buffer_load_b128(
-                    plane_rsrc, src_byte, 0, SLC));
-        } else {
+        if (z_start + dz < wnz)
+            regs[dz] = hipPlaneLoadNT<bcopy_float4_vec>(
+                d_src + (size_t)(z_start + dz) * wav_plane, src_byte);
+        else
             regs[dz] = zero_vec_from;
-        }
     }
 
-    uint32_t dst_byte = ((y0 + gy) * ldimx + x0 + gx) * (uint32_t)sizeof(float);
+    size_t dst_byte = ((size_t)(y0 + gy) * ldimx + x0 + gx) * sizeof(float);
     bool x_all_in = (gx + 3 < ex);
 
     if (x_all_in) {
         #pragma unroll
         for (int dz = 0; dz < BCOPY_ZPB; ++dz) {
             int iz = z_start + dz;
-            if (iz < ez) {
-                auto plane_rsrc = __builtin_amdgcn_make_buffer_rsrc(
-                    d_dst + (long)(z0 + iz) * ldimxy,
-                    0, -1, 0x00027000);
-                auto vi = __builtin_bit_cast(bcopy_int4_vec, regs[dz]);
-                __builtin_amdgcn_raw_buffer_store_b128(
-                    vi, plane_rsrc, dst_byte, 0, SLC);
-            }
+            if (iz < ez)
+                hipPlaneStoreF4u(
+                    d_dst + (size_t)(z0 + iz) * ldimxy, dst_byte,
+                    __builtin_bit_cast(hip_float4_u, regs[dz]));
         }
     } else {
+        // Ragged x edge: store lane-by-lane, guarded by the extraction window
+        // (global stores have no OOB clamp).
         #pragma unroll
         for (int dz = 0; dz < BCOPY_ZPB; ++dz) {
             int iz = z_start + dz;
             if (iz < ez) {
-                auto plane_rsrc = __builtin_amdgcn_make_buffer_rsrc(
-                    d_dst + (long)(z0 + iz) * ldimxy,
-                    0, -1, 0x00027000);
-                float e0 = regs[dz][0], e1 = regs[dz][1];
-                float e2 = regs[dz][2], e3 = regs[dz][3];
-                if (gx + 0 < ex) __builtin_amdgcn_raw_buffer_store_b32(
-                    __builtin_bit_cast(int, e0), plane_rsrc, dst_byte + 0, 0, SLC);
-                if (gx + 1 < ex) __builtin_amdgcn_raw_buffer_store_b32(
-                    __builtin_bit_cast(int, e1), plane_rsrc, dst_byte + 4, 0, SLC);
-                if (gx + 2 < ex) __builtin_amdgcn_raw_buffer_store_b32(
-                    __builtin_bit_cast(int, e2), plane_rsrc, dst_byte + 8, 0, SLC);
-                if (gx + 3 < ex) __builtin_amdgcn_raw_buffer_store_b32(
-                    __builtin_bit_cast(int, e3), plane_rsrc, dst_byte + 12, 0, SLC);
+                float* plane = d_dst + (size_t)(z0 + iz) * ldimxy;
+                if (gx + 0 < ex) hipPlaneStoreScalarNT(plane, dst_byte + 0,  regs[dz][0]);
+                if (gx + 1 < ex) hipPlaneStoreScalarNT(plane, dst_byte + 4,  regs[dz][1]);
+                if (gx + 2 < ex) hipPlaneStoreScalarNT(plane, dst_byte + 8,  regs[dz][2]);
+                if (gx + 3 < ex) hipPlaneStoreScalarNT(plane, dst_byte + 12, regs[dz][3]);
             }
         }
     }