src: cpu: aarch64: Enable matmul static quantisation.

Author: renato-arantes

src/cpu/aarch64/matmul/acl_lowp_matmul_sq.cpp

@@ -0,0 +1,249 @@
+/*******************************************************************************
+* Copyright 2024 Arm Ltd. and affiliates
+*
+* Licensed 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 "cpu/aarch64/matmul/acl_lowp_matmul_sq.hpp"
+
+namespace dnnl {
+namespace impl {
+namespace cpu {
+namespace aarch64 {
+namespace matmul {
+
+status_t acl_lowp_matmul_sq_resource_t::configure(
+        const acl_lowp_matmul_sq_conf_t &almc) {
+
+    if (!acl_obj_) return status::out_of_memory;
+
+    acl_obj_->src_tensor.allocator()->init(almc.src_tensor_info);
+    acl_obj_->wei_tensor.allocator()->init(almc.wei_tensor_info);
+    if (almc.with_bias) {
+        acl_obj_->bia_tensor.allocator()->init(almc.bia_tensor_info);
+    }
+    acl_obj_->dst_tensor.allocator()->init(almc.dst_tensor_info);
+
+    arm_compute::QuantizationInfo qi {1.0, 0, true};
+    acl_obj_->src_tensor.info()->set_quantization_info(qi);
+    acl_obj_->wei_tensor.info()->set_quantization_info(qi);
+    acl_obj_->dst_tensor.info()->set_quantization_info(qi);
+
+    acl_obj_->gemm.configure(&acl_obj_->src_tensor, &acl_obj_->wei_tensor,
+            almc.with_bias ? &acl_obj_->bia_tensor : nullptr,
+            &acl_obj_->dst_tensor, almc.gemm_info);
+
+    return status::success;
+}
+
+status_t acl_lowp_matmul_sq_t::pd_t::init(engine_t *engine) {
+    VDISPATCH_MATMUL(set_default_formats(), "failed to set default formats");
+    using smask_t = primitive_attr_t::skip_mask_t;
+    VDISPATCH_MATMUL(
+            attr()->has_default_values(smask_t::scales_runtime
+                    | smask_t::zero_points_runtime | smask_t::post_ops),
+            "only scale, zero point and post-ops attrs supported");
+
+    VDISPATCH_MATMUL(attr()->scales_.get(DNNL_ARG_SRC).mask_ == 0
+                    && attr()->zero_points_.get(DNNL_ARG_SRC) == 0
+                    && attr()->scales_.get(DNNL_ARG_WEIGHTS).mask_ == 0
+                    && attr()->zero_points_.get(DNNL_ARG_WEIGHTS) == 0
+                    && attr()->scales_.get(DNNL_ARG_DST).mask_ == 0
+                    && attr()->zero_points_.get(DNNL_ARG_DST) == 0,
+            "common scales and zero points only");
+
+    VDISPATCH_MATMUL(
+            !has_runtime_dims_or_strides(), VERBOSE_RUNTIMEDIM_UNSUPPORTED);
+
+    const memory_desc_wrapper src_d(src_md_);
+    const memory_desc_wrapper wei_d(weights_md_);
+    const memory_desc_wrapper bia_d(bias_md_);
+    const memory_desc_wrapper dst_d(dst_md_);
+
+    using namespace data_type;
+    VDISPATCH_MATMUL(utils::one_of(src_d.data_type(), s8, u8)
+                            && wei_d.data_type() == s8
+                            && src_d.data_type() == s8
+                    ? dst_d.data_type() == s8
+                    : dst_d.data_type() == u8
+                            && utils::one_of(bia_d.data_type(), f32, undef),
+            VERBOSE_UNSUPPORTED_DT_CFG);
+
+    VDISPATCH_MATMUL(src_d.matches_tag(format_tag::ab)
+                    && wei_d.matches_tag(format_tag::ab)
+                    && dst_d.matches_tag(format_tag::ab),
+            VERBOSE_UNSUPPORTED_TAG);
+
+    VDISPATCH_MATMUL_SC(
+            memory_desc_init_by_tag(bias_md_, bias_md_.ndims, bias_md_.dims,
+                    bias_md_.data_type, format_tag::ab),
+            VERBOSE_UNSUPPORTED_BIAS_CFG);
+
+    // We set the QuantizationInfo to be dynamic because it is re-set in run()
+    almc_.src_tensor_info
+            = arm_compute::TensorInfo(arm_compute::TensorShape(K(), M()), 1,
+                    acl_utils::get_acl_data_t(src_d.data_type(), true),
+                    arm_compute::QuantizationInfo(1.0, 0, true));
+    almc_.src_tensor_info.set_are_values_constant(false);
+
+    almc_.wei_tensor_info
+            = arm_compute::TensorInfo(arm_compute::TensorShape(N(), K()), 1,
+                    acl_utils::get_acl_data_t(wei_d.data_type(), true),
+                    arm_compute::QuantizationInfo(1.0, 0, true));
+    almc_.wei_tensor_info.set_are_values_constant(false);
+
+    almc_.dst_tensor_info
+            = arm_compute::TensorInfo(arm_compute::TensorShape(N(), M()), 1,
+                    acl_utils::get_acl_data_t(dst_d.data_type(), true),
+                    arm_compute::QuantizationInfo(1.0, 0, true));
+
+    almc_.bia_tensor_info = arm_compute::TensorInfo(
+            arm_compute::TensorShape(), 1, arm_compute::DataType::S32);
+    almc_.with_bias = bia_d.format_kind() != format_kind::undef;
+    if (almc_.with_bias) {
+        // This is not currently guarded in ACL
+        VDISPATCH_MATMUL(bia_d.ndims() == 2 && bia_d.dims()[0] == 1
+                        && bia_d.dims()[1] == N(),
+                "Only 1xN bias is supported");
+        almc_.bia_tensor_info.set_tensor_shape(
+                arm_compute::TensorShape(bia_d.dims()[1], bia_d.dims()[0]));
+    }
+
+    arm_compute::GEMMLowpOutputStageInfo info;
+
+    info.type = arm_compute::GEMMLowpOutputStageType::QUANTIZE_DOWN_FIXEDPOINT;
+    info.gemmlowp_multiplier = 1073741824;
+    info.gemmlowp_shift = -1;
+    info.gemmlowp_offset = 0;
+    info.gemmlowp_min_bound = -128;
+    info.gemmlowp_max_bound = 127;
+    info.output_data_type = almc_.dst_tensor_info.data_type();
+
+    almc_.gemm_info.set_gemmlowp_output_stage(info);
+
+    auto scratchpad = scratchpad_registry().registrar();
+    const dnnl::impl::memory_desc_t dst_md_ {desc_.dst_desc};
+    arm_compute::ActivationLayerInfo act_info;
+    CHECK(init_scratchpad(engine, scratchpad, acl_post_ops, attr_.post_ops_,
+            act_info, dst_md_));
+
+    almc_.gemm_info.set_activation_info(act_info);
+    ACL_CHECK_VALID(arm_compute::NEGEMMLowpMatrixMultiplyCore::validate(
+            &almc_.src_tensor_info, &almc_.wei_tensor_info,
+            almc_.with_bias ? &almc_.bia_tensor_info : nullptr,
+            &almc_.dst_tensor_info, almc_.gemm_info));
+
+    return status::success;
+}
+
+status_t acl_lowp_matmul_sq_t::pd_t::init_scratchpad(engine_t *engine,
+        memory_tracking::registrar_t &scratchpad, acl_post_ops_t &post_ops,
+        dnnl::impl::post_ops_t &attr_post_ops,
+        arm_compute::ActivationLayerInfo &act_info,
+        const dnnl::impl::memory_desc_t &dst_md) {
+    CHECK(post_ops.init(engine, attr_post_ops, dst_md, act_info));
+    // ACL only accepts s32 bias for quantization and since
+    // the current bias vector is f32 we need to convert.
+    if (almc_.with_bias) {
+        const memory_desc_wrapper bias_d(&bias_md_);
+        // printf("*** init_scratchpad: bias_data_type=%d\n", bias_md_.data_type);
+        scratchpad.book(memory_tracking::names::key_conv_bias_s32_convert,
+                bias_d.nelems(), bias_d.data_type_size());
+    }
+    return status::success;
+}
+
+status_t acl_lowp_matmul_sq_t::create_resource(
+        engine_t *engine, resource_mapper_t &mapper) const {
+
+    if (mapper.has_resource(this)) return status::success;
+
+    auto r = utils::make_unique<acl_lowp_matmul_sq_resource_t>();
+    if (!r) return status::out_of_memory;
+
+    CHECK(r->configure(pd()->almc_));
+
+    mapper.add(this, std::move(r));
+
+    return status::success;
+}
+
+status_t acl_lowp_matmul_sq_t::execute(const exec_ctx_t &ctx) const {
+    std::lock_guard<std::mutex> _lock {this->mtx};
+
+    bool with_bias = pd()->almc_.with_bias;
+
+    acl_lowp_matmul_sq_obj_t &acl_obj
+            = ctx.get_resource_mapper()
+                      ->get<acl_lowp_matmul_sq_resource_t>(this)
+                      ->get_acl_obj();
+
+    auto src = CTX_IN_MEM(const int8_t *, DNNL_ARG_SRC);
+    auto wei = CTX_IN_MEM(const int8_t *, DNNL_ARG_WEIGHTS);
+    auto dst = CTX_OUT_MEM(const int8_t *, DNNL_ARG_DST);
+
+    acl_obj.src_tensor.allocator()->import_memory(const_cast<int8_t *>(src));
+    acl_obj.wei_tensor.allocator()->import_memory(const_cast<int8_t *>(wei));
+    acl_obj.dst_tensor.allocator()->import_memory(const_cast<int8_t *>(dst));
+
+    DEFINE_ARG_SCALES_BUFFER(src_scale, DNNL_ARG_SRC);
+    DEFINE_ZERO_POINT_VALUE(src_zero_point, DNNL_ARG_SRC);
+    DEFINE_ARG_SCALES_BUFFER(wei_scale, DNNL_ARG_WEIGHTS);
+    DEFINE_ZERO_POINT_VALUE(wei_zero_point, DNNL_ARG_WEIGHTS);
+    DEFINE_ARG_SCALES_BUFFER(dst_scale, DNNL_ARG_DST);
+    DEFINE_ZERO_POINT_VALUE(dst_zero_point, DNNL_ARG_DST);
+
+    if (with_bias) {
+        const auto scratchpad = ctx.get_scratchpad_grantor();
+        auto bia_s32_base = scratchpad.get<uint32_t>(
+                memory_tracking::names::key_conv_bias_s32_convert);
+        auto bia_f32_base = CTX_IN_MEM(const float32_t *, DNNL_ARG_BIAS);
+        const float bias_scale = 1 / (*src_scale * (*wei_scale));
+        const int num_elements
+                = acl_obj.bia_tensor.info()->total_size() / sizeof(float32_t);
+        parallel_nd(num_elements, [&](dim_t e) {
+            const auto b = int32_t(std::round(bia_f32_base[e] * bias_scale));
+            bia_s32_base[e] = b;
+        });
+        acl_obj.bia_tensor.allocator()->init(*acl_obj.bia_tensor.info());
+        acl_obj.bia_tensor.allocator()->import_memory(bia_s32_base);
+    }
+
+    acl_obj.src_tensor.info()->set_quantization_info(
+            arm_compute::QuantizationInfo(*src_scale, -src_zero_point, true));
+    acl_obj.wei_tensor.info()->set_quantization_info(
+            arm_compute::QuantizationInfo(*wei_scale, -wei_zero_point, true));
+
+    // for efficiency reasons, OneDNN saves the inverse of the destination
+    acl_obj.dst_tensor.info()->set_quantization_info(
+            arm_compute::QuantizationInfo(
+                    1.0 / (*dst_scale), dst_zero_point, true));
+
+    acl_obj.gemm.update_quantization_parameters();
+
+    acl_obj.gemm.run();
+
+    // free() here tells ACL it can no longer use it, it does not deallocate
+    acl_obj.src_tensor.allocator()->free();
+    acl_obj.wei_tensor.allocator()->free();
+    if (with_bias) { acl_obj.bia_tensor.allocator()->free(); }
+    acl_obj.dst_tensor.allocator()->free();
+
+    return status::success;
+};
+
+} // namespace matmul
+} // namespace aarch64
+} // namespace cpu
+} // namespace impl
+} // namespace dnnl