Quantization Framework Implementation with 1bit and MultiBit Binary Q…

…uantizer (#1929)

* Quantization Framework Implementation with 1bit and MultiBit Binary Quantizer

Signed-off-by: VIKASH TIWARI <[email protected]>

* Quantization Framework Implementation with 1bit and MultiBit Binary Quantizer

Signed-off-by: VIKASH TIWARI <[email protected]>

* Implemented Serlization using Writable

Signed-off-by: VIKASH TIWARI <[email protected]>

---------

Signed-off-by: VIKASH TIWARI <[email protected]>
Signed-off-by: Vikasht34 <[email protected]>

CHANGELOG.md

@@ -35,3 +35,4 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 * Generalize lib interface to return context objects [#1925](https://github.com/opensearch-project/k-NN/pull/1925)
 * Move k search k-NN query to re-write phase of vector search query for Native Engines [#1877](https://github.com/opensearch-project/k-NN/pull/1877)
 * Restructure mappers to better handle null cases and avoid branching in parsing [#1939](https://github.com/opensearch-project/k-NN/pull/1939)
+* Added Quantization Framework and implemented 1Bit and multibit quantizer[#1889](https://github.com/opensearch-project/k-NN/issues/1889)

src/main/java/org/opensearch/knn/quantization/enums/ScalarQuantizationType.java

@@ -0,0 +1,62 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.enums;
+
+import lombok.Getter;
+
+/**
+ * The ScalarQuantizationType enum defines the various scalar quantization types that can be used
+ * for vector quantization. Each type corresponds to a different bit-width representation of the quantized values.
+ *
+ * <p>
+ * Future Developers: If you change the name of any enum constant, do not change its associated value.
+ * Serialization and deserialization depend on these values to maintain compatibility.
+ * </p>
+ */
+@Getter
+public enum ScalarQuantizationType {
+    /**
+     * ONE_BIT quantization uses a single bit per coordinate.
+     */
+    ONE_BIT(1),
+
+    /**
+     * TWO_BIT quantization uses two bits per coordinate.
+     */
+    TWO_BIT(2),
+
+    /**
+     * FOUR_BIT quantization uses four bits per coordinate.
+     */
+    FOUR_BIT(4);
+
+    private final int id;
+
+    /**
+     * Constructs a ScalarQuantizationType with the specified ID.
+     *
+     * @param id the ID representing the quantization type.
+     */
+    ScalarQuantizationType(int id) {
+        this.id = id;
+    }
+
+    /**
+     * Returns the ScalarQuantizationType associated with the given ID.
+     *
+     * @param id the ID of the quantization type.
+     * @return the corresponding ScalarQuantizationType.
+     * @throws IllegalArgumentException if the ID does not correspond to any ScalarQuantizationType.
+     */
+    public static ScalarQuantizationType fromId(int id) {
+        for (ScalarQuantizationType type : ScalarQuantizationType.values()) {
+            if (type.getId() == id) {
+                return type;
+            }
+        }
+        throw new IllegalArgumentException("Unknown ScalarQuantizationType ID: " + id);
+    }
+}

src/main/java/org/opensearch/knn/quantization/factory/QuantizerFactory.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
+import org.opensearch.knn.quantization.quantizer.Quantizer;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * The QuantizerFactory class is responsible for creating instances of {@link Quantizer}
+ * based on the provided {@link QuantizationParams}. It uses a registry to look up the
+ * appropriate quantizer implementation for the given quantization parameters.
+ */
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+public final class QuantizerFactory {
+    private static final AtomicBoolean isRegistered = new AtomicBoolean(false);
+
+    /**
+     * Ensures that default quantizers are registered.
+     */
+    private static void ensureRegistered() {
+        if (!isRegistered.get()) {
+            synchronized (QuantizerFactory.class) {
+                if (!isRegistered.get()) {
+                    QuantizerRegistrar.registerDefaultQuantizers();
+                    isRegistered.set(true);
+                }
+            }
+        }
+    }
+
+    /**
+     * Retrieves a quantizer instance based on the provided quantization parameters.
+     *
+     * @param params the quantization parameters used to determine the appropriate quantizer
+     * @param <P>    the type of quantization parameters, extending {@link QuantizationParams}
+     * @param <Q>    the type of the quantized output
+     * @return an instance of {@link Quantizer} corresponding to the provided parameters
+     */
+    public static <P extends QuantizationParams, Q> Quantizer<P, Q> getQuantizer(final P params) {
+        if (params == null) {
+            throw new IllegalArgumentException("Quantization parameters must not be null.");
+        }
+        // Lazy Registration instead of static block as class level;
+        ensureRegistered();
+        return QuantizerRegistry.getQuantizer(params);
+    }
+}

src/main/java/org/opensearch/knn/quantization/factory/QuantizerRegistrar.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.opensearch.knn.quantization.enums.ScalarQuantizationType;
+import org.opensearch.knn.quantization.models.quantizationParams.ScalarQuantizationParams;
+import org.opensearch.knn.quantization.quantizer.MultiBitScalarQuantizer;
+import org.opensearch.knn.quantization.quantizer.OneBitScalarQuantizer;
+
+/**
+ * The QuantizerRegistrar class is responsible for registering default quantizers.
+ * This class ensures that the registration happens only once in a thread-safe manner.
+ */
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+final class QuantizerRegistrar {
+
+    /**
+     * Registers default quantizers
+     * <p>
+     * This method is synchronized to ensure that registration occurs only once,
+     * even in a multi-threaded environment.
+     * </p>
+     */
+    static synchronized void registerDefaultQuantizers() {
+        // Register OneBitScalarQuantizer for SQParams with VALUE_QUANTIZATION and SQTypes.ONE_BIT
+        QuantizerRegistry.register(
+            ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.ONE_BIT),
+            new OneBitScalarQuantizer()
+        );
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 2
+        QuantizerRegistry.register(
+            ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.TWO_BIT),
+            new MultiBitScalarQuantizer(2)
+        );
+        // Register MultiBitScalarQuantizer for SQParams with VALUE_QUANTIZATION with bit per co-ordinate = 4
+        QuantizerRegistry.register(
+            ScalarQuantizationParams.generateTypeIdentifier(ScalarQuantizationType.FOUR_BIT),
+            new MultiBitScalarQuantizer(4)
+        );
+    }
+}

src/main/java/org/opensearch/knn/quantization/factory/QuantizerRegistry.java

@@ -0,0 +1,59 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.factory;
+
+import lombok.AccessLevel;
+import lombok.NoArgsConstructor;
+import org.opensearch.knn.quantization.models.quantizationParams.QuantizationParams;
+import org.opensearch.knn.quantization.quantizer.Quantizer;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * The QuantizerRegistry class is responsible for managing the registration and retrieval
+ * of quantizer instances. Quantizers are registered with specific quantization parameters
+ * and type identifiers, allowing for efficient lookup and instantiation.
+ */
+@NoArgsConstructor(access = AccessLevel.PRIVATE)
+final class QuantizerRegistry {
+    // ConcurrentHashMap for thread-safe access
+    private static final Map<String, Quantizer<?, ?>> registry = new ConcurrentHashMap<>();
+
+    /**
+     * Registers a quantizer with the registry.
+     *
+     * @param paramIdentifier the unique identifier for the quantization parameters
+     * @param quantizer       an instance of the quantizer
+     */
+    static void register(final String paramIdentifier, final Quantizer<?, ?> quantizer) {
+        // Check if the quantizer is already registered for the given identifier
+        if (registry.putIfAbsent(paramIdentifier, quantizer) != null) {
+            // Throw an exception if a quantizer is already registered
+            throw new IllegalArgumentException("Quantizer already registered for identifier: " + paramIdentifier);
+        }
+    }
+
+    /**
+     * Retrieves a quantizer instance based on the provided quantization parameters.
+     *
+     * @param params the quantization parameters used to determine the appropriate quantizer
+     * @param <P>    the type of quantization parameters
+     * @param <Q>    the type of the quantized output
+     * @return an instance of {@link Quantizer} corresponding to the provided parameters
+     * @throws IllegalArgumentException if no quantizer is registered for the given parameters
+     */
+    static <P extends QuantizationParams, Q> Quantizer<P, Q> getQuantizer(final P params) {
+        String identifier = params.getTypeIdentifier();
+        Quantizer<?, ?> quantizer = registry.get(identifier);
+        if (quantizer == null) {
+            throw new IllegalArgumentException("No quantizer registered for type identifier: " + identifier);
+        }
+        @SuppressWarnings("unchecked")
+        Quantizer<P, Q> typedQuantizer = (Quantizer<P, Q>) quantizer;
+        return typedQuantizer;
+    }
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/BinaryQuantizationOutput.java

@@ -0,0 +1,67 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationOutput;
+
+import lombok.Getter;
+import lombok.NoArgsConstructor;
+
+import java.util.Arrays;
+
+/**
+ * The BinaryQuantizationOutput class represents the output of a quantization process in binary format.
+ * It implements the QuantizationOutput interface to handle byte arrays specifically.
+ */
+@NoArgsConstructor
+public class BinaryQuantizationOutput implements QuantizationOutput<byte[]> {
+    @Getter
+    private byte[] quantizedVector;
+
+    /**
+     * Prepares the quantized vector array based on the provided parameters and returns it for direct modification.
+     * This method ensures that the internal byte array is appropriately sized and cleared before being used.
+     * The method accepts two parameters:
+     * <ul>
+     *     <li><b>bitsPerCoordinate:</b> The number of bits used per coordinate. This determines the granularity of the quantization.</li>
+     *     <li><b>vectorLength:</b> The length of the original vector that needs to be quantized. This helps in calculating the required byte array size.</li>
+     * </ul>
+     * If the existing quantized vector is either null or not the same size as the required byte array,
+     * a new byte array is allocated. Otherwise, the existing array is cleared (i.e., all bytes are set to zero).
+     * This method is designed to be used in conjunction with a bit-packing utility that writes quantized values directly
+     * into the returned byte array.
+     * @param params an array of parameters, where the first parameter is the number of bits per coordinate (int),
+     *               and the second parameter is the length of the vector (int).
+     * @return the prepared and writable quantized vector as a byte array.
+     * @throws IllegalArgumentException if the parameters are not as expected (e.g., missing or not integers).
+     */
+    @Override
+    public byte[] prepareAndGetWritableQuantizedVector(Object... params) {
+        if (params.length != 2 || !(params[0] instanceof Integer) || !(params[1] instanceof Integer)) {
+            throw new IllegalArgumentException("Expected two integer parameters: bitsPerCoordinate and vectorLength");
+        }
+        int bitsPerCoordinate = (int) params[0];
+        int vectorLength = (int) params[1];
+        int totalBits = bitsPerCoordinate * vectorLength;
+        int byteLength = (totalBits + 7) >> 3;
+
+        if (this.quantizedVector == null || this.quantizedVector.length != byteLength) {
+            this.quantizedVector = new byte[byteLength];
+        } else {
+            Arrays.fill(this.quantizedVector, (byte) 0);
+        }
+
+        return this.quantizedVector;
+    }
+
+    /**
+     * Returns the quantized vector.
+     *
+     * @return the quantized vector byte array.
+     */
+    @Override
+    public byte[] getQuantizedVector() {
+        return quantizedVector;
+    }
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationOutput/QuantizationOutput.java

@@ -0,0 +1,28 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationOutput;
+
+/**
+ * The QuantizationOutput interface defines the contract for quantization output data.
+ *
+ * @param <T> The type of the quantized data.
+ */
+public interface QuantizationOutput<T> {
+    /**
+     * Returns the quantized vector.
+     *
+     * @return the quantized data.
+     */
+    T getQuantizedVector();
+
+    /**
+     * Prepares and returns the writable quantized vector for direct modification.
+     *
+     * @param params the parameters needed for preparing the quantized vector.
+     * @return the prepared and writable quantized vector.
+     */
+    T prepareAndGetWritableQuantizedVector(Object... params);
+}

src/main/java/org/opensearch/knn/quantization/models/quantizationParams/QuantizationParams.java

@@ -0,0 +1,27 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.knn.quantization.models.quantizationParams;
+
+import org.opensearch.core.common.io.stream.Writeable;
+
+/**
+ * Interface for quantization parameters.
+ * This interface defines the basic contract for all quantization parameter types.
+ * It provides methods to retrieve the quantization type and a unique type identifier.
+ * Implementations of this interface are expected to provide specific configurations
+ * for various quantization strategies.
+ */
+public interface QuantizationParams extends Writeable {
+    /**
+     * Provides a unique identifier for the quantization parameters.
+     * This identifier is typically a combination of the quantization type
+     * and additional specifics, and it serves to distinguish between different
+     * configurations or modes of quantization.
+     *
+     * @return a string representing the unique type identifier.
+     */
+    String getTypeIdentifier();
+}