[XPU] super big ernie support (#7184)

Author: mayang002

lite/api/paddle_api.cc

@@ -547,6 +547,16 @@ void CxxConfig::set_xpu_l3_cache_method(size_t l3_size, bool locked) {
 #endif
 }
 
+void set_xpu_gm_workspace_method(size_t gm_size) {
+#ifdef LITE_WITH_XPU
+  lite::TargetWrapperXPU::local_gm_size = gm_size;
+#else
+  LOG(WARNING) << "The invoking of the function "
+                  "'set_xpu_gm_workspace_method' is ignored, please "
+                  "rebuild it with LITE_WITH_XPU=ON.";
+#endif
+}
+
 void CxxConfig::set_xpu_dev_per_thread(int dev_no) {
 #ifdef LITE_WITH_XPU
   lite::TargetWrapperXPU::SetDev(dev_no);

lite/api/paddle_api.h

@@ -408,6 +408,8 @@ class LITE_API CxxConfig : public ConfigBase {
   void set_xpu_workspace_l3_size_per_thread(int l3_size = 0x4000000);
   void set_xpu_l3_cache_method(size_t l3_size, bool locked = false);
 
+  void set_xpu_gm_workspace_method(size_t gm_size);
+
   void set_xpu_conv_autotune(bool autotune = true,
                              const std::string& autotune_file = "");
 

lite/backends/xpu/target_wrapper.cc

@@ -180,6 +180,8 @@ LITE_THREAD_LOCAL std::string TargetWrapperXPU::conv_autotune_file;
 LITE_THREAD_LOCAL bool TargetWrapperXPU::need_l3_mutex{false};
 LITE_THREAD_LOCAL size_t TargetWrapperXPU::local_l3_size{
     std::numeric_limits<size_t>::max()};
+LITE_THREAD_LOCAL size_t TargetWrapperXPU::local_gm_size{
+    0x4000000};  // 64 * 1024 * 1024
 LITE_THREAD_LOCAL void* TargetWrapperXPU::local_l3_ptr_{nullptr};
 void* TargetWrapperXPU::shared_l3_ptr_{nullptr};
 size_t TargetWrapperXPU::shared_l3_size{0};

lite/backends/xpu/target_wrapper.h

@@ -100,6 +100,22 @@ class TargetWrapper<TARGET(kXPU)> {
         local_l3_size = max_l3_size;
       }
       CHECK_LE(shared_l3_size, max_l3_size);
+      if (local_gm_size > 0) {
+        VLOG(3) << "Try To Malloc Local GM Workspace Size is" << local_gm_size;
+        void* local_gm_ptr = nullptr;
+        int ret =
+            xpu_malloc(reinterpret_cast<void**>(&local_gm_ptr), local_gm_size);
+        if (ret != 0) {
+          VLOG(3) << "No Enough GM Workspace For Current Predictor.";
+        } else {
+          ret = tls_raw_ctx_->_gm_mgr.set(local_gm_ptr, local_gm_size);
+          if (ret != 0) {
+            LOG(WARNING) << "XPU GM Mgr Init Fail, Please Check Configuration.";
+            XPU_CALL(xpu_free(local_gm_ptr));
+            local_gm_ptr = nullptr;
+          }
+        }
+      }
     }
     return tls_raw_ctx_;
   }
@@ -131,7 +147,8 @@ class TargetWrapper<TARGET(kXPU)> {
   // l3 cache config
   static LITE_THREAD_LOCAL bool need_l3_mutex;    // model level l3 size
   static LITE_THREAD_LOCAL size_t local_l3_size;  // model level l3 size
-  static size_t shared_l3_size;                   // model level l3 size
+  static LITE_THREAD_LOCAL size_t local_gm_size;
+  static size_t shared_l3_size;  // model level l3 size
   static LITE_THREAD_LOCAL std::vector<XPUL3CacheBlock*>
       l3_block_dict;  // l3 cache block used between op layers
 

lite/core/optimizer/mir/fusion/__xpu__multi_encoder_fuse_pass.cc

@@ -1046,7 +1046,7 @@ class XPUMultiEncoderFuser {
                                                    weight_qkv_trans_int8.get(),
                                                    max_f,
                                                    qkv_len);
-        memcpy(weight_tensor_vec[0]->mutable_data<float>(),
+        memcpy(weight_tensor_vec[0]->mutable_data<int8_t>(),
                weight_qkv_trans_int8.get(),
                qkv_len * sizeof(int8_t));
       } else {
@@ -1056,7 +1056,7 @@ class XPUMultiEncoderFuser {
             weight_qkv_trans_int16.get(),
             max_f,
             qkv_len);
-        memcpy(weight_tensor_vec[0]->mutable_data<float>(),
+        memcpy(weight_tensor_vec[0]->mutable_data<int16_t>(),
                weight_qkv_trans_int16.get(),
                qkv_len * sizeof(int16_t));
       }