Support jit execution of aggregate functions to improve performance of aggregate operator

CMakeLists.txt

@@ -458,6 +458,16 @@ if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64")
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fsigned-char")
 endif()
 
+# Keep frame pointers and avoid sibling-call optimization so perf can fully
+# unwind stacks with the cheap frame-pointer call-graph (no DWARF needed).
+# Off by default; enable with -DBOLT_ENABLE_FRAME_POINTER=ON when profiling.
+option(BOLT_ENABLE_FRAME_POINTER
+  "Preserve frame pointers for perf stack unwinding" OFF)
+if(BOLT_ENABLE_FRAME_POINTER)
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -fno-optimize-sibling-calls")
+  set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -fno-omit-frame-pointer -mno-omit-leaf-frame-pointer -fno-optimize-sibling-calls")
+endif()
+
 # Under Ninja, we are able to designate certain targets large enough to require restricted
 # parallelism.
 if("${MAX_HIGH_MEM_JOBS}")

Makefile

@@ -83,6 +83,11 @@ BUILD_TYPE=Release
 BOLT_BUILD_BENCHMARKS ?= "OFF"
 # Control whether to build tests with coverage instrumentation
 BOLT_BUILD_TESTING_WITH_COVERAGE ?= "OFF"
+# Control whether to keep frame pointers (for perf stack unwinding)
+BOLT_ENABLE_FRAME_POINTER ?= "OFF"
+# Control whether to report JIT symbols to Intel VTune (jitprofiling)
+BOLT_ENABLE_VTUNE_JIT ?= "OFF"
+VTUNE_SDK_DIR ?=
 # -----------------------------------------------------------------
 
 # TODO: remove `BUILD_USER` and `BUILD_CHANNEL`
@@ -204,6 +209,9 @@ conan_build: conan_install
 	NUM_THREADS=$(NUM_THREADS) \
 	BOLT_BUILD_BENCHMARKS=${BOLT_BUILD_BENCHMARKS} \
 	BOLT_BUILD_TESTING_WITH_COVERAGE=${BOLT_BUILD_TESTING_WITH_COVERAGE} \
+	BOLT_ENABLE_FRAME_POINTER=${BOLT_ENABLE_FRAME_POINTER} \
+	BOLT_ENABLE_VTUNE_JIT=${BOLT_ENABLE_VTUNE_JIT} \
+	VTUNE_SDK_DIR=${VTUNE_SDK_DIR} \
 	conan build ../.. --name=bolt --version=${BUILD_VERSION} --user=${BUILD_USER} --channel=${BUILD_CHANNEL} \
 	   -s llvm-core/*:build_type=Release \
 	   -s "&:build_type=${BUILD_TYPE}" \
@@ -306,6 +314,17 @@ benchmarks-build-spark:
 benchmarks-build-relwithdebinfo:
 	$(MAKE) conan_build BUILD_TYPE=RelWithDebInfo BOLT_BUILD_BENCHMARKS="ON" CONAN_CONFIG=" -c bolt/*:tools.build:skip_test=False" CONAN_OPTIONS="-o bolt/*:spark_compatible=False -o bolt/*:enable_testutil=True -o bolt/*:enable_perf=True"
 
+# Same as benchmarks-build-spark but keeps frame pointers so perf can unwind
+# stacks with the cheap frame-pointer call-graph (no DWARF needed).
+benchmarks-build-spark-profile:
+	$(MAKE) conan_build BUILD_TYPE=Release BOLT_BUILD_BENCHMARKS="ON" BOLT_ENABLE_FRAME_POINTER="ON" CONAN_CONFIG=" -c bolt/*:tools.build:skip_test=False" CONAN_OPTIONS="-o bolt/*:spark_compatible=True -o bolt/*:enable_testutil=True -o bolt/*:enable_perf=True"
+
+# Same as benchmarks-build-spark-profile but also reports JIT symbols to Intel
+# VTune (libjitprofiling). Override the SDK path with VTUNE_SDK_DIR=... if it is
+# not under the default /opt/intel/oneapi/vtune/2023.2.0/sdk.
+benchmarks-build-spark-vtune:
+	$(MAKE) conan_build BUILD_TYPE=Release BOLT_BUILD_BENCHMARKS="ON" BOLT_ENABLE_FRAME_POINTER="ON" BOLT_ENABLE_VTUNE_JIT="ON" VTUNE_SDK_DIR="${VTUNE_SDK_DIR}" CONAN_CONFIG=" -c bolt/*:tools.build:skip_test=False" CONAN_OPTIONS="-o bolt/*:spark_compatible=True -o bolt/*:enable_testutil=True -o bolt/*:enable_perf=True"
+
 unittest_debug: unittest
 unittest: debug_with_test
 	ctest --test-dir $(BUILD_BASE_DIR)/Debug --timeout 7200 -j $(NUM_THREADS) --output-on-failure

bolt/common/base/AggregationStats.h

@@ -28,5 +28,12 @@ struct AggregationStats {
   uint64_t aggOutputTimeNs{0};
   uint64_t aggProbeBypassTimeNs{0};
   uint64_t aggProbeBypassCount{0};
+  // Hash aggregation JIT fine-grained timing.
+  // One-time codegen (LLVM compile) time for the JIT plan.
+  uint64_t aggJitCodegenTimeNs{0};
+  // JIT-executed part of the agg function update time.
+  uint64_t aggFunctionJitTimeNs{0};
+  // JIT-executed part of the extracting groups time.
+  uint64_t aggExtractGroupsJitTimeNs{0};
 };
 } // namespace bytedance::bolt::common

bolt/core/QueryConfig.h

@@ -652,6 +652,12 @@ class QueryConfig {
    */
   static constexpr const char* kJitLevel = "jit.level";
 
+  static constexpr const char* kHashAggrJitEnabled = "hashaggr.jit.enabled";
+  static constexpr const char* kHashAggrJitMinFuseWidth =
+      "hashaggr.jit.min_fuse_width";
+  static constexpr const char* kHashAggrJitMaxFuseWidth =
+      "hashaggr.jit.max_fuse_width";
+
   // expired, to deleted later
   static constexpr const char* kBoltJitEnabled = "bolt.jit.enabled";
   // For morsel-driven Bolt
@@ -1606,6 +1612,18 @@ class QueryConfig {
     return flag & 1;
   }
 
+  bool enableHashAggrJit() const {
+    return get<bool>(kHashAggrJitEnabled, true);
+  }
+
+  int32_t hashAggrJitMinFuseWidth() const {
+    return get<int32_t>(kHashAggrJitMinFuseWidth, 1);
+  }
+
+  int32_t hashAggrJitMaxFuseWidth() const {
+    return get<int32_t>(kHashAggrJitMaxFuseWidth, 16);
+  }
+
   int exceptionTraceLevel() const {
     return get<int>(kExceptionTraceLevel, 1);
   }

bolt/exec/Aggregate.cpp

@@ -321,4 +321,27 @@ void Aggregate::clearInternal() {
   numNulls_ = 0;
 }
 
+#ifdef ENABLE_BOLT_JIT
+bool Aggregate::supportsHashAggrJit(
+    const jit::HashAggrJitPlanContext& /*context*/) const {
+  return false;
+}
+
+std::optional<jit::HashAggrJitDescriptor> Aggregate::createHashAggrJitDescriptor(
+    const jit::HashAggrJitPlanContext& /*context*/) const {
+  return std::nullopt;
+}
+
+jit::HashAggrJitSlot Aggregate::createHashAggrJitSlot(
+    int32_t aggregateIndex,
+    const jit::HashAggrJitDescriptor& descriptor) const {
+  return jit::HashAggrJitSlot{
+      .aggregateIndex = aggregateIndex,
+      .offset = accumulatorOffset(),
+      .nullByte = accumulatorNullByte(),
+      .nullMask = accumulatorNullMask(),
+      .desc = descriptor};
+}
+#endif
+
 } // namespace bytedance::bolt::exec

bolt/exec/Aggregate.h

@@ -33,12 +33,17 @@
 #include <folly/CPortability.h>
 #include <folly/Synchronized.h>
 
+#include <optional>
+
 #include "bolt/common/memory/HashStringAllocator.h"
 #include "bolt/core/PlanNode.h"
 #include "bolt/core/QueryConfig.h"
 #include "bolt/exec/AggregateUtil.h"
 #include "bolt/expression/FunctionSignature.h"
 #include "bolt/functions/InlineFlatten.h"
+#ifdef ENABLE_BOLT_JIT
+#include "bolt/jit/aggregation/HashAggrJitTypes.h"
+#endif
 #include "bolt/vector/BaseVector.h"
 namespace bytedance::bolt::core {
 class ExpressionEvaluator;
@@ -66,6 +71,18 @@ class Aggregate {
     return resultType_;
   }
 
+  int32_t accumulatorOffset() const {
+    return offset_;
+  }
+
+  int32_t accumulatorNullByte() const {
+    return nullByte_;
+  }
+
+  uint8_t accumulatorNullMask() const {
+    return nullMask_;
+  }
+
   // Returns the fixed number of bytes the accumulator takes on a group
   // row. Variable width accumulators will reference the variable
   // width part of the state from the fixed part.
@@ -100,6 +117,18 @@ class Aggregate {
     return false;
   }
 
+#ifdef ENABLE_BOLT_JIT
+  virtual bool supportsHashAggrJit(
+      const jit::HashAggrJitPlanContext& context) const;
+
+  virtual std::optional<jit::HashAggrJitDescriptor> createHashAggrJitDescriptor(
+      const jit::HashAggrJitPlanContext& context) const;
+
+  jit::HashAggrJitSlot createHashAggrJitSlot(
+      int32_t aggregateIndex,
+      const jit::HashAggrJitDescriptor& descriptor) const;
+#endif
+
   void setAllocator(HashStringAllocator* allocator) {
     setAllocatorInternal(allocator);
     pool_ = allocator->pool();
@@ -133,6 +162,10 @@ class Aggregate {
     setOffsetsInternal(offset, nullByte, nullMask, rowSizeOffset);
   }
 
+  void markNullCountUnknown() {
+    numNulls_ = std::nullopt;
+  }
+
   // Initializes null flags and accumulators for newly encountered groups.  This
   // function should be called only once for each group.
   //
@@ -360,7 +393,15 @@ class Aggregate {
   }
 
   bool isNull(char* group) const {
-    return numNulls_ && (group[nullByte_] & nullMask_);
+    return mayHaveNulls() && (group[nullByte_] & nullMask_);
+  }
+
+  bool hasNoNulls() const {
+    return numNulls_.has_value() && *numNulls_ == 0;
+  }
+
+  bool mayHaveNulls() const {
+    return !numNulls_.has_value() || *numNulls_ > 0;
   }
 
   // Sets null flag for all specified groups to true.
@@ -369,26 +410,33 @@ class Aggregate {
     for (auto i : indices) {
       groups[i][nullByte_] |= nullMask_;
     }
-    numNulls_ += indices.size();
+    if (numNulls_.has_value()) {
+      *numNulls_ += indices.size();
+    }
   }
 
   inline bool setNull(char* group) {
     if (group[nullByte_] & nullMask_) {
       return false;
     }
     group[nullByte_] |= nullMask_;
-    ++numNulls_;
+    if (numNulls_.has_value()) {
+      ++*numNulls_;
+    }
     return true;
   }
 
   inline bool clearNull(char* group) {
-    if (numNulls_) {
-      uint8_t mask = group[nullByte_];
-      if (mask & nullMask_) {
-        group[nullByte_] = mask & ~nullMask_;
-        --numNulls_;
-        return true;
+    if (!mayHaveNulls()) {
+      return false;
+    }
+    uint8_t mask = group[nullByte_];
+    if (mask & nullMask_) {
+      group[nullByte_] = mask & ~nullMask_;
+      if (numNulls_.has_value()) {
+        --*numNulls_;
       }
+      return true;
     }
     return false;
   }
@@ -449,9 +497,11 @@ class Aggregate {
   int32_t rowSizeOffset_ = 0;
 
   // Number of null accumulators in the current state of the aggregation
-  // operator for this aggregate. If 0, clearing the null as part of update
-  // is not needed.
-  uint64_t numNulls_ = 0;
+  // operator for this aggregate.
+  // - 0          => known that no group is null
+  // - N > 0      => known exact null count
+  // - nullopt    => unknown; must rely on per-group null bit
+  std::optional<uint64_t> numNulls_{0};
   HashStringAllocator* allocator_{nullptr};
   memory::MemoryPool* pool_{nullptr};
   std::shared_ptr<core::ExpressionEvaluator> expressionEvaluator_{nullptr};

bolt/exec/AggregateCompanionAdapter.cpp

@@ -35,13 +35,15 @@
 #include "bolt/exec/RowContainer.h"
 #include "bolt/expression/SignatureBinder.h"
 #include "bolt/functions/lib/aggregates/AggregateToIntermediate.h"
+
 namespace bytedance::bolt::exec {
 
 void AggregateCompanionFunctionBase::setOffsetsInternal(
     int32_t offset,
     int32_t nullByte,
     uint8_t nullMask,
     int32_t rowSizeOffset) {
+  Aggregate::setOffsetsInternal(offset, nullByte, nullMask, rowSizeOffset);
   fn_->setOffsets(offset, nullByte, nullMask, rowSizeOffset);
 }
 
@@ -65,6 +67,19 @@ bool AggregateCompanionFunctionBase::supportsToIntermediate() const {
   return fn_->supportsToIntermediate();
 }
 
+#ifdef ENABLE_BOLT_JIT
+bool AggregateCompanionFunctionBase::supportsHashAggrJit(
+    const jit::HashAggrJitPlanContext& context) const {
+  return fn_->supportsHashAggrJit(rewriteHashAggrJitContext(context));
+}
+
+std::optional<jit::HashAggrJitDescriptor>
+AggregateCompanionFunctionBase::createHashAggrJitDescriptor(
+    const jit::HashAggrJitPlanContext& context) const {
+  return fn_->createHashAggrJitDescriptor(rewriteHashAggrJitContext(context));
+}
+#endif
+
 bool AggregateCompanionFunctionBase::supportAccumulatorSerde() const {
   return fn_->supportAccumulatorSerde();
 }