cpu: aarch64: hot fix for aux tensor management of stateless gemm and winograd conv

Signed-off-by: Ye Tao <ye.tao@arm.com>

Author: taoye9

src/cpu/aarch64/acl_gemm_convolution.cpp

@@ -98,11 +98,26 @@ status_t acl_gemm_convolution_fwd_t<src_t, wei_t, dst_t, bia_t>::init(
     return status::success;
 }
 
+template <data_type_t src_t, data_type_t wei_t, data_type_t dst_t,
+        data_type_t bia_t>
+void acl_gemm_convolution_fwd_t<src_t, wei_t, dst_t,
+        bia_t>::reinitialize_acl_obj() const {
+    auto acp = pd()->acp_;
+    std::lock_guard<std::mutex> _lock {this->mtx};
+    acl_obj_ = std::make_unique<acl_obj_t<Op>>();
+    acl_obj_->conv.configure(&acp.src_tensor_info, &acp.wei_tensor_info,
+            acp.with_bias ? &acp.bia_tensor_info : nullptr,
+            &acp.dst_tensor_info, acp.padstride_info, acp.weights_info,
+            acp.dilation_info, acp.act_info, acp.fast_math);
+    acl_obj_->aux_mem_req = acl_obj_->conv.workspace();
+}
+
 template <data_type_t src_t, data_type_t wei_t, data_type_t dst_t,
         data_type_t bia_t>
 status_t
 acl_gemm_convolution_fwd_t<src_t, wei_t, dst_t, bia_t>::execute_forward(
         const exec_ctx_t &ctx) const {
+    reinitialize_acl_obj();
     return execute_forward_conv_acl<acl_obj_t<Op>, pd_t, src_data_t, wei_data_t,
             dst_data_t, bia_data_t>(ctx, acl_obj_.get(), pd(), gemm_conv_keys);
 }

src/cpu/aarch64/acl_gemm_convolution.hpp

@@ -65,8 +65,14 @@ struct acl_gemm_convolution_fwd_t : public primitive_t {
 
 private:
     status_t execute_forward(const exec_ctx_t &ctx) const;
+
+    // hot fix solution for stateless API which should be replaced soon.
+    mutable std::mutex mtx;
+    void reinitialize_acl_obj() const;
+
     const pd_t *pd() const { return (const pd_t *)primitive_t::pd().get(); }
-    std::unique_ptr<acl_obj_t<Op>> acl_obj_;
+    mutable std::unique_ptr<acl_obj_t<Op>> acl_obj_;
+
 }; // acl_gemm_convolution_fwd_t
 
 } // namespace aarch64

src/cpu/aarch64/acl_winograd_convolution.cpp

@@ -129,17 +129,21 @@ status_t acl_wino_convolution_fwd_t::pd_t::init_conf() {
     return status::success;
 }
 
-status_t acl_wino_convolution_fwd_t::execute_forward(
-        const exec_ctx_t &ctx) const {
-    acl_obj_ = std::make_unique<acl_obj_t<Op>>();
+void acl_wino_convolution_fwd_t::reinitialize_acl_obj() const {
     auto acp = pd()->acp_;
+    std::lock_guard<std::mutex> _lock {this->mtx};
+    acl_obj_ = std::make_unique<acl_obj_t<Op>>();
     acl_obj_->conv.configure(&acp.src_tensor_info, &acp.wei_tensor_info,
             acp.with_bias ? &acp.bia_tensor_info : nullptr,
             &acp.dst_tensor_info, acp.padstride_info, acp.act_info,
             true); // to support 5x5, 7x7 filter shapes in addition to 3x3
 
     acl_obj_->aux_mem_req = acl_obj_->conv.workspace();
+}
 
+status_t acl_wino_convolution_fwd_t::execute_forward(
+        const exec_ctx_t &ctx) const {
+    reinitialize_acl_obj();
     return execute_forward_conv_acl<acl_obj_t<Op>, pd_t, data_t>(
             ctx, acl_obj_.get(), pd(), wino_conv_keys);
 }

src/cpu/aarch64/acl_winograd_convolution.hpp

@@ -58,6 +58,11 @@ struct acl_wino_convolution_fwd_t : public primitive_t {
 
 private:
     status_t execute_forward(const exec_ctx_t &ctx) const;
+
+    // hot fix solution for stateless API which should be replaced soon.
+    mutable std::mutex mtx;
+    void reinitialize_acl_obj() const;
+
     const pd_t *pd() const { return (const pd_t *)primitive_t::pd().get(); }
     mutable std::unique_ptr<acl_obj_t<Op>> acl_obj_;
 }; // acl_wino_convolution_fwd_t