[WIP][POC] Pfor encoding

cpp/cmake_modules/SetupCxxFlags.cmake

@@ -633,21 +633,12 @@ endif()
 
 if(NOT MSVC)
   set(C_RELEASE_FLAGS "")
-  if(CMAKE_C_FLAGS_RELEASE MATCHES "-O3")
-    string(APPEND C_RELEASE_FLAGS " -O2")
-  endif()
+  # Local override: keep -O3 from CMake's default Release flags for the
+  # pfor benchmark — the bench is sensitive to inlining/unrolling that
+  # -O3 enables. Upstream Arrow downgrades to -O2 here; we skip that.
   set(CXX_RELEASE_FLAGS "")
-  if(CMAKE_CXX_FLAGS_RELEASE MATCHES "-O3")
-    string(APPEND CXX_RELEASE_FLAGS " -O2")
-  endif()
   set(C_RELWITHDEBINFO_FLAGS "")
-  if(CMAKE_C_FLAGS_RELWITHDEBINFO MATCHES "-O3")
-    string(APPEND C_RELWITHDEBINFO_FLAGS " -O2")
-  endif()
   set(CXX_RELWITHDEBINFO_FLAGS "")
-  if(CMAKE_CXX_FLAGS_RELWITHDEBINFO MATCHES "-O3")
-    string(APPEND CXX_RELWITHDEBINFO_FLAGS " -O2")
-  endif()
   if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
     string(APPEND C_RELEASE_FLAGS " -ftree-vectorize")
     string(APPEND CXX_RELEASE_FLAGS " -ftree-vectorize")

cpp/src/arrow/util/CMakeLists.txt

@@ -118,6 +118,22 @@ add_arrow_test(threading-utility-test
                test_common.cc
                thread_pool_test.cc)
 
+add_arrow_test(pfor-test
+               SOURCES
+               pfor/pfor_test.cc
+               pfor/pfor.cc
+               pfor/pfor_wrapper.cc)
+
+add_arrow_benchmark(pfor/pfor_benchmark
+                    EXTRA_SOURCES
+                    pfor/pfor.cc
+                    pfor/pfor_wrapper.cc)
+
+add_arrow_test(fastlanes-for-test
+               SOURCES
+               fastlanes/fastlanes_for_test.cc
+               fastlanes/fastlanes_for.cc)
+
 add_arrow_benchmark(bit_block_counter_benchmark)
 add_arrow_benchmark(bit_util_benchmark)
 add_arrow_benchmark(bitmap_reader_benchmark)

cpp/src/arrow/util/fastlanes/fastlanes_for.cc

@@ -0,0 +1,245 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/util/fastlanes/fastlanes_for.h"
+
+#include <algorithm>
+#include <cstring>
+#include <limits>
+
+#include "arrow/util/fastlanes/fastlanes_kernels.h"
+
+namespace arrow {
+namespace util {
+namespace fastlanes {
+
+namespace {
+
+// Dispatch by runtime bit_width into the templated PackBlock<W>.
+inline void PackBlockDispatch(uint32_t bit_width, const uint32_t* in,
+                              uint32_t* out) {
+  switch (bit_width) {
+#define CASE_W(N) \
+  case N:         \
+    PackBlock<N>(in, out); \
+    break
+    CASE_W(1);  CASE_W(2);  CASE_W(3);  CASE_W(4);  CASE_W(5);  CASE_W(6);
+    CASE_W(7);  CASE_W(8);  CASE_W(9);  CASE_W(10); CASE_W(11); CASE_W(12);
+    CASE_W(13); CASE_W(14); CASE_W(15); CASE_W(16); CASE_W(17); CASE_W(18);
+    CASE_W(19); CASE_W(20); CASE_W(21); CASE_W(22); CASE_W(23); CASE_W(24);
+    CASE_W(25); CASE_W(26); CASE_W(27); CASE_W(28); CASE_W(29); CASE_W(30);
+    CASE_W(31); CASE_W(32);
+#undef CASE_W
+    default: break;
+  }
+}
+
+inline void UnpackBlockDispatch(uint32_t bit_width, const uint32_t* packed,
+                                uint32_t* out) {
+  switch (bit_width) {
+#define CASE_W(N) \
+  case N:         \
+    UnpackBlock<N>(packed, out); \
+    break
+    CASE_W(1);  CASE_W(2);  CASE_W(3);  CASE_W(4);  CASE_W(5);  CASE_W(6);
+    CASE_W(7);  CASE_W(8);  CASE_W(9);  CASE_W(10); CASE_W(11); CASE_W(12);
+    CASE_W(13); CASE_W(14); CASE_W(15); CASE_W(16); CASE_W(17); CASE_W(18);
+    CASE_W(19); CASE_W(20); CASE_W(21); CASE_W(22); CASE_W(23); CASE_W(24);
+    CASE_W(25); CASE_W(26); CASE_W(27); CASE_W(28); CASE_W(29); CASE_W(30);
+    CASE_W(31); CASE_W(32);
+#undef CASE_W
+    default: break;
+  }
+}
+
+inline uint8_t BitsRequired(uint32_t range) {
+  if (range == 0) return 0;
+  return static_cast<uint8_t>(32 - __builtin_clz(range));
+}
+
+}  // namespace
+
+Result<int64_t> FastLanesForCodec::Encode(const int32_t* input,
+                                          int64_t num_values, uint8_t* output) {
+  if (num_values % kChunkSize != 0) {
+    return Status::Invalid("FastLanesForCodec::Encode: num_values (", num_values,
+                           ") must be a multiple of ", kChunkSize);
+  }
+  const int64_t num_chunks = num_values / kChunkSize;
+  uint8_t* out_ptr = output;
+
+  // Scratch buffers: one chunk's worth of FOR-shifted values (uint32_t) in
+  // transposed FL_ORDER, ready to feed PackBlock twice (one per 1024-block).
+  alignas(64) uint32_t transposed[kChunkSize];
+
+  for (int64_t chunk = 0; chunk < num_chunks; ++chunk) {
+    const int32_t* in_chunk = input + chunk * kChunkSize;
+
+    // Compute min over the 2048 values.
+    int32_t min_val = std::numeric_limits<int32_t>::max();
+    int32_t max_val = std::numeric_limits<int32_t>::min();
+    for (int64_t i = 0; i < kChunkSize; ++i) {
+      const int32_t v = in_chunk[i];
+      if (v < min_val) min_val = v;
+      if (v > max_val) max_val = v;
+    }
+    const uint32_t range = static_cast<uint32_t>(max_val - min_val);
+    const uint8_t bit_width = BitsRequired(range);
+
+    // Write 5-byte header: [min (4B LE int32_t)] [bit_width (1B)].
+    std::memcpy(out_ptr, &min_val, sizeof(int32_t));
+    out_ptr[4] = bit_width;
+    out_ptr += kChunkHeaderSize;
+
+    if (bit_width == 0) {
+      // All values equal min; no payload bytes needed.
+      continue;
+    }
+
+    // Build the transposed FOR-shifted scratch: per 1024-block, gather
+    // input[fromTransposed32(t)] - min into transposed[t].
+    for (int64_t block = 0; block < 2; ++block) {
+      const int32_t* block_in = in_chunk + block * static_cast<int64_t>(kBlockSize);
+      uint32_t* block_t = transposed + block * kBlockSize;
+      for (size_t t = 0; t < kBlockSize; ++t) {
+        block_t[t] = static_cast<uint32_t>(block_in[fromTransposed32(t)] - min_val);
+      }
+    }
+
+    // Pack each 1024-block; payload is 2 * 128 * bit_width bytes total.
+    const int64_t bytes_per_block = 128 * bit_width;
+    for (int64_t block = 0; block < 2; ++block) {
+      PackBlockDispatch(bit_width, transposed + block * kBlockSize,
+                        reinterpret_cast<uint32_t*>(out_ptr + block * bytes_per_block));
+    }
+    out_ptr += 2 * bytes_per_block;
+  }
+
+  return out_ptr - output;
+}
+
+Status FastLanesForCodec::Decode(int32_t* output, int64_t num_values,
+                                  const uint8_t* input, int64_t input_size) {
+  if (num_values % kChunkSize != 0) {
+    return Status::Invalid("FastLanesForCodec::Decode: num_values (", num_values,
+                           ") must be a multiple of ", kChunkSize);
+  }
+  const int64_t num_chunks = num_values / kChunkSize;
+  const uint8_t* in_ptr = input;
+  const uint8_t* in_end = input + input_size;
+
+  alignas(64) uint32_t transposed[kChunkSize];
+
+  for (int64_t chunk = 0; chunk < num_chunks; ++chunk) {
+    if (in_ptr + kChunkHeaderSize > in_end) {
+      return Status::Invalid("FastLanesForCodec::Decode: truncated input header");
+    }
+    int32_t min_val;
+    std::memcpy(&min_val, in_ptr, sizeof(int32_t));
+    const uint8_t bit_width = in_ptr[4];
+    in_ptr += kChunkHeaderSize;
+
+    int32_t* out_chunk = output + chunk * kChunkSize;
+
+    if (bit_width == 0) {
+      // All values equal min; emit min in transposed order (just fill).
+      for (int64_t i = 0; i < kChunkSize; ++i) out_chunk[i] = min_val;
+      continue;
+    }
+
+    const int64_t bytes_per_block = 128 * bit_width;
+    if (in_ptr + 2 * bytes_per_block > in_end) {
+      return Status::Invalid("FastLanesForCodec::Decode: truncated input body");
+    }
+
+    // Unpack each 1024-block into the transposed scratch.
+    for (int64_t block = 0; block < 2; ++block) {
+      UnpackBlockDispatch(
+          bit_width,
+          reinterpret_cast<const uint32_t*>(in_ptr + block * bytes_per_block),
+          transposed + block * kBlockSize);
+    }
+    in_ptr += 2 * bytes_per_block;
+
+    // Add min back. Output stays in transposed FL_ORDER (no scatter).
+    for (int64_t i = 0; i < kChunkSize; ++i) {
+      out_chunk[i] = static_cast<int32_t>(transposed[i]) + min_val;
+    }
+  }
+
+  return Status::OK();
+}
+
+Status FastLanesForCodec::DecodeFlat(int32_t* output, int64_t num_values,
+                                      const uint8_t* input, int64_t input_size) {
+  if (num_values % kChunkSize != 0) {
+    return Status::Invalid("FastLanesForCodec::DecodeFlat: num_values (",
+                           num_values, ") must be a multiple of ", kChunkSize);
+  }
+  const int64_t num_chunks = num_values / kChunkSize;
+  const uint8_t* in_ptr = input;
+  const uint8_t* in_end = input + input_size;
+
+  alignas(64) uint32_t transposed[kChunkSize];
+
+  for (int64_t chunk = 0; chunk < num_chunks; ++chunk) {
+    if (in_ptr + kChunkHeaderSize > in_end) {
+      return Status::Invalid("FastLanesForCodec::DecodeFlat: truncated header");
+    }
+    int32_t min_val;
+    std::memcpy(&min_val, in_ptr, sizeof(int32_t));
+    const uint8_t bit_width = in_ptr[4];
+    in_ptr += kChunkHeaderSize;
+
+    int32_t* out_chunk = output + chunk * kChunkSize;
+
+    if (bit_width == 0) {
+      for (int64_t i = 0; i < kChunkSize; ++i) out_chunk[i] = min_val;
+      continue;
+    }
+
+    const int64_t bytes_per_block = 128 * bit_width;
+    if (in_ptr + 2 * bytes_per_block > in_end) {
+      return Status::Invalid("FastLanesForCodec::DecodeFlat: truncated body");
+    }
+
+    for (int64_t block = 0; block < 2; ++block) {
+      UnpackBlockDispatch(
+          bit_width,
+          reinterpret_cast<const uint32_t*>(in_ptr + block * bytes_per_block),
+          transposed + block * kBlockSize);
+    }
+    in_ptr += 2 * bytes_per_block;
+
+    // Scatter via fromTransposed32 per 1024-block to restore original
+    // input order: out_chunk[block*1024 + fromTransposed32(t)] = transposed +
+    // min. This is the FL_ORDER inverse of the encode-time gather.
+    for (int64_t block = 0; block < 2; ++block) {
+      const uint32_t* block_t = transposed + block * kBlockSize;
+      int32_t* block_out = out_chunk + block * static_cast<int64_t>(kBlockSize);
+      for (size_t t = 0; t < kBlockSize; ++t) {
+        block_out[fromTransposed32(t)] = static_cast<int32_t>(block_t[t]) + min_val;
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace fastlanes
+}  // namespace util
+}  // namespace arrow

cpp/src/arrow/util/fastlanes/fastlanes_for.h

@@ -0,0 +1,83 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// FastLanes + Frame-of-Reference encoder/decoder.
+//
+// Splits the input into 2048-value chunks. Per chunk:
+//   1. Compute min over the 2048 values.
+//   2. Subtract min from every value -> all values in [0, max-min].
+//   3. Compute bit_width = ceil(log2(max-min+1)).
+//   4. Apply the FastLanes FL_ORDER permutation within each 1024-block
+//      (gather input[fromTransposed32(t)]) and pack with FastLanes
+//      lane-interleaved bit-packing at bit_width.
+//
+// The decoded output is in the FastLanes transposed order (NO scatter
+// step on decode), i.e. output[chunk*2048 + block*1024 + t] equals
+// input[chunk*2048 + block*1024 + fromTransposed32(t)] + min.
+//
+// Per-chunk header (5 bytes): [min (4B int32_t little-endian)] [bit_width (1B)]
+// followed by 2 * 128 * bit_width bytes of packed data (zero if bit_width=0).
+
+#pragma once
+
+#include <cstdint>
+
+#include "arrow/result.h"
+#include "arrow/status.h"
+
+namespace arrow {
+namespace util {
+namespace fastlanes {
+
+class FastLanesForCodec {
+ public:
+  // Number of values per encoded chunk (= 2 FastLanes 1024-blocks).
+  static constexpr int64_t kChunkSize = 2048;
+
+  // Per-chunk header size: 4 bytes (int32 min) + 1 byte (bit_width).
+  static constexpr int64_t kChunkHeaderSize = 5;
+
+  // Upper-bound size of the encoded output for `num_values` (must be a
+  // multiple of kChunkSize).
+  static int64_t MaxEncodedSize(int64_t num_values) {
+    const int64_t num_chunks = (num_values + kChunkSize - 1) / kChunkSize;
+    // Max bit_width = 32 → 2 * 128 * 32 = 8192 bytes per chunk packed body.
+    return num_chunks * (kChunkHeaderSize + 2 * 128 * 32);
+  }
+
+  // Encode `num_values` int32_t values. `num_values` must be a multiple
+  // of kChunkSize. Returns the number of bytes written to `output`.
+  // `output` must be at least MaxEncodedSize(num_values) bytes.
+  static Result<int64_t> Encode(const int32_t* input, int64_t num_values,
+                                uint8_t* output);
+
+  // Decode `num_values` int32_t values. Output is in FastLanes
+  // transposed order within each 1024-block (see header comment).
+  static Status Decode(int32_t* output, int64_t num_values,
+                       const uint8_t* input, int64_t input_size);
+
+  // Decode and scatter back to ORIGINAL input order: output[i] == input[i]
+  // for the encoded input. Pays the FL_ORDER scatter cost on top of the
+  // kernel — slower than Decode(), included for apples-to-apples
+  // comparison against codecs that produce flat output.
+  static Status DecodeFlat(int32_t* output, int64_t num_values,
+                           const uint8_t* input, int64_t input_size);
+};
+
+}  // namespace fastlanes
+}  // namespace util
+}  // namespace arrow