Mdr lossless fixes

include/mgard-x/Lossless/ParallelHuffman/Huffman.hpp

@@ -204,14 +204,24 @@ class Huffman {
     // mark("Huffman stage: codebook");
 
     if (target_cr > 1.0) {
-      workspace.freq_array.hostCopy(false, queue_idx);
-      workspace.CL_array.hostCopy(false, queue_idx);
+      // Encoded bits = sum over symbols of frequency x codeword length. After
+      // GetCodebook, freq_subarray holds the frequencies sorted ascending and
+      // CL_subarray has been reversed by GenerateCW, so they cannot be paired
+      // by index. Use the unsorted histogram (_d_freq_copy_subarray) and the
+      // symbol-indexed codebook, whose top byte is the codeword length (see
+      // deflate_bitwidth).
+      std::vector<unsigned int> freq(dict_size);
+      std::vector<H> codebook(dict_size);
+      MemoryManager<DeviceType>::Copy1D(freq.data(),
+                                        workspace._d_freq_copy_subarray.data(),
+                                        dict_size, queue_idx);
+      MemoryManager<DeviceType>::Copy1D(codebook.data(),
+                                        workspace.codebook_subarray.data(),
+                                        dict_size, queue_idx);
       DeviceRuntime<DeviceType>::SyncQueue(queue_idx);
-      unsigned int *_freq = workspace.freq_array.dataHost();
-      unsigned int *_cl = workspace.CL_array.dataHost();
       double LC = 0;
       for (SIZE i = 0; i < dict_size; i++) {
-        LC += (double)_freq[i] * _cl[i];
+        LC += (double)freq[i] * (double)(codebook[i] >> (sizeof(H) * 8 - 8));
       }
       double estimated_cr =
           (double)(sizeof(Q) * primary_count) / (LC / 8 + 2000);

include/mgard-x/MDR-X/LosslessCompressor/HybridLevelCompressor.hpp

@@ -7,6 +7,7 @@
 // #include "../RefactorUtils.hpp"
 #include "LevelCompressorInterface.hpp"
 #include "LosslessCompressor.hpp"
+#include <cstring>
 
 namespace mgard_x {
 namespace MDR {
@@ -65,7 +66,7 @@ class HybridLevelCompressor
                  int level_idx, int queue_idx) {
 
     std::vector<float> cr, time;
-    bool huffman_success, rle_success;
+    bool huffman_success, rle_success, zstd_success;
     for (SIZE bitplane_idx = 0; bitplane_idx < encoded_bitplanes.shape(0);
          bitplane_idx++) {
       if (bitplane_idx % num_merged_bitplanes == 0) {
@@ -81,8 +82,14 @@ class HybridLevelCompressor
         log::level = 0;
         huffman_success = false;
         rle_success = false;
+        zstd_success = false;
         // cr_threshold = 2.0;
-        if (merged_bitplane_size > size_threshold) {
+        if (merged_bitplane_size > size_threshold &&
+            config.lossless == lossless_type::Huffman_Zstd) {
+          zstd_success =
+              compress_zstd((Byte *)bitplane, merged_bitplane_size,
+                            compressed_bitplanes[bitplane_idx], queue_idx);
+        } else if (merged_bitplane_size > size_threshold) {
           rle_success =
               rle.Compress(encoded_bitplane, compressed_bitplanes[bitplane_idx],
                            cr_threshold, queue_idx);
@@ -105,7 +112,7 @@ class HybridLevelCompressor
           }
         }
 
-        if (huffman_success == false && rle_success == false) {
+        if (!huffman_success && !rle_success && !zstd_success) {
           // direct copy
           compressed_bitplanes[bitplane_idx].resize({merged_bitplane_size});
           MemoryManager<DeviceType>::Copy1D(
@@ -172,6 +179,10 @@ class HybridLevelCompressor
           rle.Deserialize(compressed_bitplanes[bitplane_idx], queue_idx);
           rle.Decompress(compressed_bitplanes[bitplane_idx], encoded_bitplane,
                          queue_idx);
+        } else if (is_zstd(compressed_bitplanes[bitplane_idx],
+                           merged_bitplane_size, queue_idx)) {
+          decompress_zstd(compressed_bitplanes[bitplane_idx], (Byte *)bitplane,
+                          merged_bitplane_size, queue_idx);
         } else {
           // Direct copy
           MemoryManager<DeviceType>::Copy1D(
@@ -191,6 +202,52 @@ class HybridLevelCompressor
     // log::info("Time: " + time_string);
   }
 
+  // ZSTD stage (Config::lossless == Huffman_Zstd): replaces RLE/byte Huffman
+  // for groups above size_threshold and is kept whenever it is smaller than
+  // the raw group. Stored as [signature][Zstd stream].
+  static constexpr Byte zstd_signature[7] = {'M', 'G', 'X', 'Z', 'S', 'T', 'D'};
+
+  bool compress_zstd(Byte *group, SIZE n, Array<1, Byte, DeviceType> &out,
+                     int queue_idx) {
+    Array<1, Byte, DeviceType> buffer({n});
+    MemoryManager<DeviceType>::Copy1D(buffer.data(), group, n, queue_idx);
+    zstd.Compress(buffer, queue_idx);
+    SIZE size = buffer.shape(0);
+    if (size + sizeof(zstd_signature) >= n) {
+      return false;
+    }
+    out.resize({(SIZE)(size + sizeof(zstd_signature))}, queue_idx);
+    MemoryManager<DeviceType>::Copy1D(out.data(), (Byte *)zstd_signature,
+                                      sizeof(zstd_signature), queue_idx);
+    MemoryManager<DeviceType>::Copy1D(out.data() + sizeof(zstd_signature),
+                                      buffer.data(), size, queue_idx);
+    DeviceRuntime<DeviceType>::SyncQueue(queue_idx);
+    return true;
+  }
+
+  // A raw group is exactly n bytes; a ZSTD group is smaller and signed.
+  bool is_zstd(Array<1, Byte, DeviceType> &data, SIZE n, int queue_idx) {
+    if (data.shape(0) >= n || data.shape(0) <= sizeof(zstd_signature)) {
+      return false;
+    }
+    Byte signature[sizeof(zstd_signature)];
+    MemoryManager<DeviceType>::Copy1D(signature, data.data(),
+                                      sizeof(zstd_signature), queue_idx);
+    DeviceRuntime<DeviceType>::SyncQueue(queue_idx);
+    return std::memcmp(signature, zstd_signature, sizeof(zstd_signature)) == 0;
+  }
+
+  void decompress_zstd(Array<1, Byte, DeviceType> &data, Byte *group, SIZE n,
+                       int queue_idx) {
+    SIZE size = data.shape(0) - sizeof(zstd_signature);
+    Array<1, Byte, DeviceType> buffer({size});
+    MemoryManager<DeviceType>::Copy1D(
+        buffer.data(), data.data() + sizeof(zstd_signature), size, queue_idx);
+    zstd.Decompress(buffer, queue_idx);
+    MemoryManager<DeviceType>::Copy1D(group, buffer.data(), n, queue_idx);
+    DeviceRuntime<DeviceType>::SyncQueue(queue_idx);
+  }
+
   // release the buffer created
   void decompress_release() {}
 

include/mgard-x/MDR-X/Reconstructor/ComposedReconstructor.hpp

@@ -286,8 +286,9 @@ class ComposedReconstructor
       // This ensure all each batch of merged bitplanes are used for
       // Reconstruction. Otherwise, unsed bitplanes will not be guaranteed
       // to be in memory in future reconstructions.
+      // (a level that needs no bitplanes stays at 0).
       int m = Compressor::num_merged_bitplanes;
-      n = ((n - 1) / m + 1) * m;
+      n = (n + m - 1) / m * m;
     }
     timer.end();
     // timer.print("Preprocessing");
@@ -316,8 +317,10 @@ class ComposedReconstructor
 
   void LoadMetadata(MDRMetadata &mdr_metadata, MDRData<DeviceType> &mdr_data,
                     int queue_idx) {
-    for (int level_idx = 0; level_idx <= mdr_metadata.CurrFinalLevel();
-         level_idx++) {
+    // All levels, not just up to CurrFinalLevel(): levels with no bitplanes
+    // must get level_num_bitplanes = 0 rather than keep a value from a
+    // previous use of this reconstructor (ProgressiveReconstruct visits all).
+    for (int level_idx = 0; level_idx <= hierarchy->l_target(); level_idx++) {
       level_num_bitplanes[level_idx] =
           mdr_metadata.loaded_level_num_bitplanes[level_idx] -
           mdr_metadata.prev_used_level_num_bitplanes[level_idx];

src/mgard-x/Executables/mdr-x.cpp

@@ -43,6 +43,7 @@ void print_usage_message(std::string error) {
 \t\t (optional) -m / --max-memory <max memory usage>  \n\
 \t\t (optional) -dd / --domain-decomposition <max-dim|block>\n\
 \t\t\t (optional) -dd-size / --domain-decomposition-size <integer> (for block domain decomposition only) \n\
+\t\t (optional) -l / --lossless <huffman|huffman-zstd>: bitplane lossless stage (default: huffman)\n\
 \n\
 \t -x / --reconstruct: reconstruct data\n\
 \t\t -i / --input <path to refactored data dir>\n\
@@ -306,10 +307,12 @@ int launch_refactor(mgard_x::DIM D, enum mgard_x::data_type dtype,
                     std::vector<mgard_x::SIZE> shape,
                     std::string domain_decomposition, mgard_x::SIZE block_size,
                     enum mgard_x::device_type dev_type, int verbose,
-                    mgard_x::SIZE max_memory_footprint) {
+                    mgard_x::SIZE max_memory_footprint,
+                    enum mgard_x::lossless_type lossless) {
 
   mgard_x::Config config;
   config.normalize_coordinates = false;
+  config.lossless = lossless;
   config.log_level = verbose_to_log_level(verbose);
   config.decomposition = mgard_x::decomposition_type::MultiDim;
   if (domain_decomposition == "max-dim") {
@@ -498,8 +501,12 @@ bool try_refactoring(int argc, char *argv[]) {
   enum mgard_x::data_type dtype = get_data_type(argc, argv);
   std::vector<mgard_x::SIZE> shape =
       get_args<mgard_x::SIZE>(argc, argv, "Dimensions", "-dim", "--dimension");
-  // std::string lossless_level = get_arg<std::string>(argc, argv, "Lossless",
-  // "-l", "--lossless");
+  enum mgard_x::lossless_type lossless = mgard_x::lossless_type::Huffman;
+  if (has_arg(argc, argv, "-l", "--lossless") &&
+      get_arg<std::string>(argc, argv, "Lossless", "-l", "--lossless") ==
+          "huffman-zstd") {
+    lossless = mgard_x::lossless_type::Huffman_Zstd;
+  }
   enum mgard_x::device_type dev_type = get_device_type(argc, argv);
   int verbose = 0;
   if (has_arg(argc, argv, "-v", "--verbose")) {
@@ -524,12 +531,13 @@ bool try_refactoring(int argc, char *argv[]) {
   if (dtype == mgard_x::data_type::Double) {
     launch_refactor<double>(shape.size(), dtype, input_file.c_str(),
                             output_file.c_str(), shape, domain_decomposition,
-                            block_size, dev_type, verbose,
-                            max_memory_footprint);
+                            block_size, dev_type, verbose, max_memory_footprint,
+                            lossless);
   } else if (dtype == mgard_x::data_type::Float) {
     launch_refactor<float>(shape.size(), dtype, input_file.c_str(),
                            output_file.c_str(), shape, domain_decomposition,
-                           block_size, dev_type, verbose, max_memory_footprint);
+                           block_size, dev_type, verbose, max_memory_footprint,
+                           lossless);
   }
   return true;
 }