M,

FUSE_RELU,

NUM_THREADS,

INPUT_SCALE,

INPUT_OFFSET,

OUTPUT_SCALE,

OUTPUT_OFFSET,

SIGNED_INPUT

arm_compute::NEGEMMLowpMatrixMultiplyCore gemm;

// key for use in the cache

QuantMatmulCacheKey key;

QuantMatmul(

int64_t weight_dim_0,

int64_t weight_dim_1,

double weight_scale,

int64_t weight_offset,

int8_t* weight_ptr,

std::optional<float*> bias_ptr,

const QuantMatmulCacheKey& cache_key);

virtual ~QuantMatmul();

virtual arm_compute::Status validate() = 0;

virtual void configure() = 0;

protected:

arm_compute::Tensor wei_q_tensor_;

std::optional<arm_compute::Tensor> bia_tensor_;

arm_compute::GEMMInfo gemm_info_;

std::optional<arm_compute::ActivationLayerInfo> relu_info_;

arm_compute::Tensor src_q_tensor;

arm_compute::Tensor src_tensor;

arm_compute::Tensor dst_tensor;

arm_compute::NEQuantizationLayer quant;

// We need a ReLU layer here (unlike static quantization) because the ReLU

// cannot be "truly" fused with the GEMM through gemm_info in ACL dynamically

// quantized matmuls.

std::optional<arm_compute::NEActivationLayer> relu;

DynamicQuantMatmul(

int64_t weight_dim_0,

int64_t weight_dim_1,

double weight_scale,

int64_t weight_offset,

int8_t* weight_ptr,

std::optional<float*> bias_ptr,

const QuantMatmulCacheKey& cache_key);

~DynamicQuantMatmul() override;

arm_compute::Status validate() override;

void configure() override;

private:

at::Tensor src_q_tensor_orig_;

arm_compute::Tensor src_q_tensor;

arm_compute::Tensor dst_q_tensor;

StaticQuantMatmul(

int64_t weight_dim_0,

int64_t weight_dim_1,

double weight_scale,

int64_t weight_offset,

int8_t* weight_ptr,

std::optional<float*> bias_ptr,

const QuantMatmulCacheKey& cache_key);

~StaticQuantMatmul() override;

arm_compute::Status validate() override;

void configure() override;

private:

std::optional<arm_compute::Tensor> bia_q_tensor_;

std::optional<at::Tensor> bia_q_tensor_orig_;

arm_compute::Tensor qa_tensor;

arm_compute::Tensor qb_tensor;

arm_compute::Tensor qdst_tensor;

arm_compute::NEArithmeticAddition q_add;

QuantAdd(

arm_compute::DataType dtype,

const std::vector<int64_t>& input_dims,

double qa_scale,

int64_t qa_offset,

double qb_scale,

int64_t qb_offset,

double dst_scale,

int64_t dst_offset);

arm_compute::Status validate();

void configure();

private:

arm_compute::ConvertPolicy policy{arm_compute::ConvertPolicy::SATURATE};

using ACLQuantMatmul = at::native::acl_utils::QuantMatmul;

using ACLDynamicQuantMatmul = at::native::acl_utils::DynamicQuantMatmul;

using ACLStaticQuantMatmul = at::native::acl_utils::StaticQuantMatmul;

using ACLQuantMatmulCacheKey = at::native::acl_utils::QuantMatmulCacheKey;

using ACLQuantMatmulCacheKeyIndex =

at::native::acl_utils::QuantMatmulCacheKeyIndex;

PackedLinearWeightsACL(

std::unique_ptr<ideep::tensor> weight,

std::optional<ideep::tensor> bias,

at::Tensor orig_weight,

std::optional<at::Tensor> orig_bias);

at::Tensor apply_dynamic(at::Tensor input, bool reduce_range = false)

override;

at::Tensor apply_dynamic_relu(at::Tensor input, bool reduce_range = false)

override;

at::Tensor apply(

at::Tensor input,

double output_scale,

int64_t output_zero_point) override;

at::Tensor apply_relu(

at::Tensor input,

double output_scale,

int64_t output_zero_point) override;

template <typename ACLQuantMatmulT>

std::shared_ptr<ACLQuantMatmulT> get_acl_quant_matmul(

const ACLQuantMatmulCacheKey& key) {

return std::dynamic_pointer_cast<ACLQuantMatmulT>(

fetch_or_create_acl_quant_matmul<ACLQuantMatmulT>(key));

}

private:

int64_t k_;

int64_t n_;

int64_t weight_zero_point_;

double weight_scale_;

// A 2 element (per layer) cache. Given it's not intended to store more than 2

// elements, we do not need a fancy implementation. The idea behind it is to

// allow for a (configuration free) fast path for autoregressive

// transformer-like models which usually involve 2 input tensor shapes; one

// for the prefill phase and another for the autoregressive phase

std::array<std::shared_ptr<ACLQuantMatmul>, 2> cache_;

template <typename ACLQuantMatmulT>

std::shared_ptr<ACLQuantMatmul> fetch_or_create_acl_quant_matmul(

const ACLQuantMatmulCacheKey& key) {

// We're only maintaining a 2 element LRU cache

// hit first

if (cache_[0] != nullptr && cache_[0]->key == key) {

return cache_[0];

}

// hit second

if (cache_[1] != nullptr && cache_[1]->key == key) {

// Update LRU

std::swap(cache_[0], cache_[1]);

return cache_[0];

}

// miss -> replace Least Recently Used - i.e. element at index 1

cache_[1] = create_acl_quant_matmul<ACLQuantMatmulT>(key);

std::swap(cache_[0], cache_[1]);

return cache_[0];

}

template <typename ACLQuantMatmulT>

std::shared_ptr<ACLQuantMatmulT> create_acl_quant_matmul(

const ACLQuantMatmulCacheKey& key) {

std::optional<float*> bias_ptr;

if (bias_.has_value()) {

bias_ptr = (float*)bias_.value().get_data_handle();

}

auto acl_gemm = std::make_shared<ACLQuantMatmulT>(

k_,

n_,

weight_scale_,

weight_zero_point_,

(int8_t*)weight_.get()->get_data_handle(),

bias_ptr,

key);

// validate

auto status = acl_gemm->validate();

if (status.error_code() != arm_compute::ErrorCode::OK) {

TORCH_WARN(

"Arm Compute Library's Quantized Matmul Validation Failed: " +

status.error_description());

return nullptr;

}

// configure

acl_gemm->configure();

return acl_gemm;

}

template <bool ReluFused>

at::Tensor apply_dynamic_impl(at::Tensor input, bool reduce_range = false);

template <bool ReluFused>

at::Tensor apply_impl(

at::Tensor input,

double output_scale,

int64_t output_zero_point);