added SIMD pragma on AMS outer loop

gropaul · Jul 30, 2024 · 8df8105 · 8df8105
1 parent d1b56c2
commit 8df8105
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 48 deletions.
diff --git a/src/common/types/CMakeLists.txt b/src/common/types/CMakeLists.txt
@@ -10,7 +10,6 @@ add_library_unity(
   duckdb_common_types
   OBJECT
   ams_sketch.cpp
-  ams_sketch_simple.cpp
   batched_data_collection.cpp
   bit.cpp
   blob.cpp

diff --git a/src/common/types/ams_sketch_simple.cpp b/src/common/types/ams_sketch_simple.cpp
diff --git a/src/execution/operator/join/physical_hash_join.cpp b/src/execution/operator/join/physical_hash_join.cpp
@@ -245,12 +245,13 @@ void PhysicalHashJoin::UpdateAMS(DataChunk &chunk, JoinHashTable &ht) const {
 
 	// get the data from the payload
 	auto *chain_key_hashes = FlatVector::GetData<hash_t>(chain_key_hashes_v);
+#pragma clang loop vectorize(assume_safety)
 	for (idx_t i = 0; i < payload_size; i++) {
 		hash_t hash = chain_key_hashes[i];
 		ht.ams_sketch_simple.Update(hash);
-//		if (hash == 0xffffffffffffffff) {
-//			printf("hash is 0\n");
-//		}
+		//		if (hash == 0xffffffffffffffff) {
+		//			printf("hash is 0\n");
+		//		}
 	}
 }
 

diff --git a/src/include/duckdb/common/types/ams_sketch_simple.hpp b/src/include/duckdb/common/types/ams_sketch_simple.hpp
@@ -3,27 +3,52 @@
 
 namespace duckdb {
 
-template<uint64_t ArraySize, uint8_t NHashFunctions>
+inline uint8_t ShiftToByteIndex(uint64_t hash, uint8_t byte_index) {
+	return static_cast<uint8_t>(hash >> (byte_index * 8));
+}
+
+inline uint8_t ShiftByteToBitIndex(uint8_t byte, uint8_t bit_index) {
+	return (byte >> bit_index) & 1;
+}
+
+template<uint8_t ArraySize, uint8_t NHashFunctions>
 class AMSSketchSimple {
 public:
 
 	// as ArraySize is a power of 2, we can use a bitmask instead of modulo
-	static constexpr uint64_t ARRAY_BITMASK = ArraySize - 1;
+	static constexpr uint8_t ARRAY_BITMASK = ArraySize - 1;
+
+	int64_t flat_array[ArraySize * NHashFunctions];
+	uint64_t update_count;
 
 	explicit AMSSketchSimple() : update_count(0) {
 		// Initialize the flat_array to zero
-		for (uint64_t i = 0; i < ArraySize * NHashFunctions; i++) {
+		for (uint16_t i = 0; i < ArraySize * NHashFunctions; i++) {
 			flat_array[i] = 0;
 		}
 
 		// the array size must be a power of 2
 		static_assert((ArraySize & (ArraySize - 1)) == 0, "ArraySize must be a power of 2");
 	}
 
-	inline void Update(uint64_t hash);
+	inline void Update(uint64_t hash){
+		update_count++;
+
+		for (uint8_t hash_function_index = 0; hash_function_index < NHashFunctions; hash_function_index++) {
+			uint8_t byte = ShiftToByteIndex(hash, hash_function_index);
+			uint8_t increment = ShiftByteToBitIndex(byte, 7);
+			int8_t sign = 1 - 2 * increment;
+
+			uint8_t byte_offset = byte & this->ARRAY_BITMASK;
+
+			// Max size is 8 * 128 = 2048, so we can use uint16_t
+			uint16_t flat_index = hash_function_index * ArraySize + byte_offset;
+			flat_array[flat_index] += sign;
+		}
+	}
 
 	void Combine(AMSSketchSimple<ArraySize, NHashFunctions>& other) {
-		for (uint64_t i = 0; i < ArraySize * NHashFunctions; i++) {
+		for (uint16_t i = 0; i < ArraySize * NHashFunctions; i++) {
 			flat_array[i] += other.flat_array[i];
 		}
 		this->update_count += other.update_count;
@@ -33,16 +58,13 @@ class AMSSketchSimple {
 	vector<vector<int64_t>> GetArray() const {
 		vector<vector<int64_t>> array(NHashFunctions, vector<int64_t>(ArraySize));
 		for (uint8_t hash_function_index = 0; hash_function_index < NHashFunctions; hash_function_index++) {
-			for (uint64_t index = 0; index < ArraySize; index++) {
-				array[hash_function_index][index] = flat_array[hash_function_index * ArraySize + index];
+			for (uint8_t index = 0; index < ArraySize; index++) {
+				uint16_t flat_index = hash_function_index * ArraySize + index;
+				array[hash_function_index][index] = flat_array[flat_index];
 			}
 		}
 		return array;
 	}
-
-private:
-	int64_t flat_array[ArraySize * NHashFunctions];
-	uint64_t update_count;
 };
 
 } // namespace duckdb