text: /*******************************************************************************
  * Copyright 2022 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 tos 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.
  *******************************************************************************/
  
  #ifndef CPU_AARCH64_ACL_BINARY_HPP
  #define CPU_AARCH64_ACL_BINARY_HPP
  
  #include "cpu/cpu_binary_pd.hpp"
  
  #include "cpu/aarch64/acl_utils.hpp"
  
  namespace dnnl {
  namespace impl {
  namespace cpu {
  namespace aarch64 {
  
  struct acl_binary_obj_t {
      std::unique_ptr<arm_compute::IFunction> binary_op;
      arm_compute::Tensor src0_tensor;
      arm_compute::Tensor src1_tensor;
      arm_compute::Tensor dst_tensor;
  };
  
  struct acl_binary_conf_t {
      arm_compute::TensorInfo src0_info;
      arm_compute::TensorInfo src1_info;
      arm_compute::TensorInfo dst_info;
      alg_kind_t alg;
  };
  
  struct acl_binary_resource_t : public resource_t {
      acl_binary_resource_t()
          : acl_obj_(utils::make_unique<acl_binary_obj_t>()) {}
  
      status_t configure(const acl_binary_conf_t &asp) {
          if (!acl_obj_) return status::out_of_memory;
  
          // Init Compute Library tensors based on info from descriptor
          acl_obj_->src0_tensor.allocator()->init(asp.src0_info);
          acl_obj_->src1_tensor.allocator()->init(asp.src1_info);
          acl_obj_->dst_tensor.allocator()->init(asp.dst_info);
  
          switch (asp.alg) {
              case alg_kind::binary_add: {
                  auto add_op
                          = std::make_unique<arm_compute::NEArithmeticAddition>();
                  add_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor,
                          arm_compute::ConvertPolicy::SATURATE);
                  acl_obj_->binary_op = std::move(add_op);
                  break;
              }
              case alg_kind::binary_sub: {
                  auto sub_op = std::make_unique<
                          arm_compute::NEArithmeticSubtraction>();
                  sub_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor,
                          arm_compute::ConvertPolicy::SATURATE);
                  acl_obj_->binary_op = std::move(sub_op);
                  break;
              }
              case alg_kind::binary_div: {
                  auto div_op = std::make_unique<
                          arm_compute::NEElementwiseDivision>();
                  div_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor);
                  acl_obj_->binary_op = std::move(div_op);
                  break;
              }
              case alg_kind::binary_mul: {
                  auto mul_op = std::make_unique<
                          arm_compute::NEPixelWiseMultiplication>();
                  mul_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor, 1.0f,
                          arm_compute::ConvertPolicy::SATURATE,
                          arm_compute::RoundingPolicy::TO_ZERO);
                  acl_obj_->binary_op = std::move(mul_op);
                  break;
              }
              case alg_kind::binary_min: {
                  auto min_op = std::make_unique<arm_compute::NEElementwiseMin>();
                  min_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor);
                  acl_obj_->binary_op = std::move(min_op);
                  break;
              }
              case alg_kind::binary_max: {
                  auto max_op = std::make_unique<arm_compute::NEElementwiseMax>();
                  max_op->configure(&acl_obj_->src0_tensor,
                          &acl_obj_->src1_tensor, &acl_obj_->dst_tensor);
                  acl_obj_->binary_op = std::move(max_op);
                  break;
              }
              default: return status::runtime_error;
          }
  
          return status::success;
      }
  
      acl_binary_obj_t &get_acl_obj() const { return *acl_obj_; }
  
      DNNL_DISALLOW_COPY_AND_ASSIGN(acl_binary_resource_t);
  
  private:
      std::unique_ptr<acl_binary_obj_t> acl_obj_;
  }; // acl_binary_resource_t
  
  struct acl_binary_t : public primitive_t {
      struct pd_t : public cpu_binary_pd_t {
          using cpu_binary_pd_t::cpu_binary_pd_t;
  
          DECLARE_COMMON_PD_T("acl", acl_binary_t);
  
          status_t init(engine_t *engine) {
  
              using namespace acl_utils;
  
              // Only support f32 and s32 for now
              data_type_t ddt = dst_md(0)->data_type;
              if (ddt != data_type::f32 && ddt != data_type::s32)
                  return status::unimplemented;
  
              // Only support src and dst all matching for now
              if (ddt != src_md(0)->data_type || src_md(1)->data_type != ddt)
                  return status::unimplemented;
  
              // Sets the memory format of dst from any to src_md(0) blocking desc
              CHECK(set_default_params());
  
              if (!attr()->has_default_values()) return status::unimplemented;
  
              asp_.alg = desc()->alg_kind;
  
              // All the algorithms we support
              if (!utils::one_of(asp_.alg, alg_kind::binary_add,
                          alg_kind::binary_sub, alg_kind::binary_mul,
                          alg_kind::binary_div, alg_kind::binary_max,
                          alg_kind::binary_min))
                  return status::unimplemented;
  
              // s32 div in ACL does not round as oneDNN expects
              if (ddt == data_type::s32 && asp_.alg == alg_kind::binary_div)
                  return status::unimplemented;
  
              // ACL pointwise arithmetic operators assume that the innermost
              // dimensions are dense for src0, src1 and dst. So we try to permute
              // the logical dims so that the innermost dim on each desc is dense
              // (without any data reordering)
              memory_desc_t src_d0_permed, src_d1_permed, dst_d_permed;
              CHECK(permute_common_dense_dimension_to_last(&src_d0_permed,
                      &src_d1_permed, &dst_d_permed, src_md(0), src_md(1),
                      dst_md()));
  
              // Create ACL tensor infos with permuted descs
              CHECK(tensor_info(asp_.src0_info, src_d0_permed));
              CHECK(tensor_info(asp_.src1_info, src_d1_permed));
              CHECK(tensor_info(asp_.dst_info, dst_d_permed));
  
              // This forces ACL not to parallelise with small workloads, this is
              // a temporary fix and should be removed in future versions (TODO)
              memory_desc_wrapper dst_d(dst_md());
              if (dst_d.nelems() < 40000) {
                  size_t acl_y_axis_i = 1;
                  CHECK(insert_singleton_dimension(asp_.src0_info, acl_y_axis_i));
                  CHECK(insert_singleton_dimension(asp_.src1_info, acl_y_axis_i));
                  CHECK(insert_singleton_dimension(asp_.dst_info, acl_y_axis_i));
              }
  
              // Call operator specific validate function to check support
              ACL_CHECK_VALID(validate(asp_));
  
              return status::success;
          }
  
          acl_binary_conf_t asp_;
  
      private:
          arm_compute::Status validate(const acl_binary_conf_t &asp) {
              switch (asp.alg) {
                  case alg_kind::binary_add:
                      return arm_compute::NEArithmeticAddition::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info,
                              arm_compute::ConvertPolicy::SATURATE);
                  case alg_kind::binary_sub:
                      return arm_compute::NEArithmeticSubtraction::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info,
                              arm_compute::ConvertPolicy::SATURATE);
                  case alg_kind::binary_div:
                      return arm_compute::NEElementwiseDivision::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info);
                  case alg_kind::binary_mul:
                      return arm_compute::NEPixelWiseMultiplication::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info, 1.0f,
                              arm_compute::ConvertPolicy::SATURATE,
                              arm_compute::RoundingPolicy::TO_ZERO);
                  case alg_kind::binary_min:
                      return arm_compute::NEElementwiseMin::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info);
                  case alg_kind::binary_max:
                      return arm_compute::NEElementwiseMax::validate(
                              &asp.src0_info, &asp.src1_info, &asp.dst_info);
                  default:
                      return arm_compute::Status(
                              arm_compute::ErrorCode::RUNTIME_ERROR,
                              "unsupported alg_kind");
              }
          }
  
      }; // pd_t
  
      acl_binary_t(const pd_t *apd) : primitive_t(apd) {}
  
      status_t create_resource(
              engine_t *engine, resource_mapper_t &mapper) const override {
          if (mapper.has_resource(this)) return status::success;
  
          auto r = utils::make_unique<acl_binary_resource_t>();
          if (!r) return status::out_of_memory;
  
          // Configure the resource based on information from primitive descriptor
          auto st = r->configure(pd()->asp_);
          if (st == status::success) { mapper.add(this, std::move(r)); }
  
          return st;
      }
  
      status_t execute(const exec_ctx_t &ctx) const override {
          return execute_forward(ctx);
      }
  
  private:
      // To guard the const execute_forward, the mutex must be 'mutable'
      mutable std::mutex mtx;
      status_t execute_forward(const exec_ctx_t &ctx) const;
      const pd_t *pd() const { return (const pd_t *)primitive_t::pd().get(); }
  
  }; // acl_binary_t
  
  } // namespace aarch64
  } // namespace cpu
  } // namespace impl
  } // namespace dnnl
  
  #endif
  