[LITE][XPU] Support mmdnn3.0-ras (a.k.a. crmm-0608) (#3950)

* fix typo

* [LITE][XPU] accomodate crmm(variant 20200608)

* refine lite/tests/api/test_mmdnn_lite_xpu.cc

* more comments, test=develop test=xpu

* bugfix in crmm pattern match

* pr comments, test=develop test=xpu

* add XPU_CALL and retval check, test=develop test=xpu

lite/backends/xpu/debug.h

@@ -19,7 +19,7 @@
 #include <memory>
 #include <string>
 #include <type_traits>
-#include "lite/backends/xpu/xpu_header_sitter.h"
+#include "lite/backends/xpu/target_wrapper.h"
 
 namespace paddle {
 namespace lite {
@@ -82,8 +82,8 @@ void DumpXPUMem(const T* ptr,
                 size_t item_per_line = 30) {
   size_t after_stride_len = (len + stride - 1) / stride;
   std::unique_ptr<T[]> cpu_mem(new T[len]);
-  xpu_memcpy(
-      cpu_mem.get(), ptr, len * sizeof(T), XPUMemcpyKind::XPU_DEVICE_TO_HOST);
+  XPU_CALL(xpu_memcpy(
+      cpu_mem.get(), ptr, len * sizeof(T), XPUMemcpyKind::XPU_DEVICE_TO_HOST));
   std::unique_ptr<T[]> after_stride(new T[after_stride_len]);
   for (size_t i = 0; i < after_stride_len; ++i) {
     after_stride[i] = cpu_mem[i * stride];

lite/backends/xpu/target_wrapper.cc

@@ -19,11 +19,11 @@ namespace lite {
 
 void* TargetWrapperXPU::Malloc(size_t size) {
   void* ptr{nullptr};
-  xpu_malloc(&ptr, size);
+  XPU_CALL(xpu_malloc(&ptr, size));
   return ptr;
 }
 
-void TargetWrapperXPU::Free(void* ptr) { xpu_free(ptr); }
+void TargetWrapperXPU::Free(void* ptr) { XPU_CALL(xpu_free(ptr)); }
 
 void TargetWrapperXPU::MemcpySync(void* dst,
                                   const void* src,
@@ -31,10 +31,10 @@ void TargetWrapperXPU::MemcpySync(void* dst,
                                   IoDirection dir) {
   switch (dir) {
     case IoDirection::HtoD:
-      xpu_memcpy(dst, src, size, XPU_HOST_TO_DEVICE);
+      XPU_CALL(xpu_memcpy(dst, src, size, XPU_HOST_TO_DEVICE));
       break;
     case IoDirection::DtoH:
-      xpu_memcpy(dst, src, size, XPU_DEVICE_TO_HOST);
+      XPU_CALL(xpu_memcpy(dst, src, size, XPU_DEVICE_TO_HOST));
       break;
     default:
       LOG(FATAL) << "Unsupported IoDirection " << static_cast<int>(dir);
@@ -49,7 +49,7 @@ XPUScratchPadGuard TargetWrapperXPU::MallocScratchPad(size_t size,
   } else {
     ptr = TargetWrapperXPU::Malloc(size);
   }
-  CHECK(ptr != nullptr);
+  CHECK(ptr != nullptr) << "size = " << size << ", use_l3 = " << use_l3;
   return XPUScratchPadGuard(new XPUScratchPad(ptr, use_l3));
 }
 

lite/backends/xpu/target_wrapper.h

@@ -16,11 +16,23 @@
 
 #include <memory>                                 // std::unique_ptr
 #include "lite/backends/xpu/xpu_header_sitter.h"  // xpu_free
-#include "lite/core/target_wrapper.h"
+#include "lite/core/target_wrapper.h"             // TargetWrapper
+#include "lite/utils/cp_logging.h"                // CHECK_EQ
+
+#define XPU_CALL(func)                                        \
+  {                                                           \
+    auto e = (func);                                          \
+    CHECK_EQ(e, 0) << "XPU: (" << #func << ") returns " << e; \
+  }
 
 namespace paddle {
 namespace lite {
 
+// MAX(lod.size()) = 64
+const int XPU_MAX_LOD_SIZE = 64;
+// MAX(lod[i + 1] - lod[i]) = 512
+const int XPU_MAX_LOD_SEQ_LEN = 512;
+
 using TargetWrapperXPU = TargetWrapper<TARGET(kXPU)>;
 
 struct XPUScratchPad {
@@ -33,7 +45,7 @@ struct XPUScratchPad {
 struct XPUScratchPadDeleter {
   void operator()(XPUScratchPad* sp) const {
     if (!sp->is_l3_) {
-      xpu_free(sp->addr_);
+      XPU_CALL(xpu_free(sp->addr_));
     }
     delete sp;
   }
@@ -55,7 +67,7 @@ class TargetWrapper<TARGET(kXPU)> {
                          size_t size,
                          IoDirection dir);
 
-  static XPUScratchPadGuard MallocScratchPad(size_t size, bool use_l3 = true);
+  static XPUScratchPadGuard MallocScratchPad(size_t size, bool use_l3 = false);
 
   static xdnn::Context* GetRawContext() {
     if (tls_raw_ctx_ == nullptr) {
@@ -77,11 +89,10 @@ class TargetWrapper<TARGET(kXPU)> {
   static void SetDev(int dev_no = 0) {
     const char* dev_env = getenv("LITE_XPU_DEV");
     if (dev_env) {
-      xpu_set_device(atoi(dev_env));
-      return;
+      dev_no = atoi(dev_env);
     }
 
-    xpu_set_device(dev_no);
+    XPU_CALL(xpu_set_device(dev_no));
   }
 
   static std::string multi_encoder_precision;  // NOLINT

lite/kernels/xpu/__xpu__embedding_with_eltwise_add_compute.cc

@@ -31,11 +31,14 @@ void XPUEmbeddingWithEltwiseAddCompute::PrepareForRun() {
     CHECK_EQ(table_dims.size(), 2); /* shape like [table_len, embed_dim] */
     table_lens_cpu_.push_back(table_dims[0]);
   }
-  void* lens_ptr = nullptr;
+
   size_t lens_size = table_lens_cpu_.size() * sizeof(int);
-  xpu_malloc(&lens_ptr, lens_size);
-  xpu_memcpy(lens_ptr, &table_lens_cpu_[0], lens_size, XPU_HOST_TO_DEVICE);
-  table_lens_guard_.reset(lens_ptr);
+  table_lens_guard_ =
+      TargetWrapperXPU::MallocScratchPad(lens_size, false /* use_l3 */);
+  XPU_CALL(xpu_memcpy(table_lens_guard_->addr_,
+                      &table_lens_cpu_[0],
+                      lens_size,
+                      XPU_HOST_TO_DEVICE));
 }
 
 void XPUEmbeddingWithEltwiseAddCompute::Run() {
@@ -55,16 +58,16 @@ void XPUEmbeddingWithEltwiseAddCompute::Run() {
   int embed_dim = table_dims[1];
   int emb_layer_num = param.Ids.size();
   int r = xdnn::embedding_with_ewadd<float, int64_t, false, false>(
-      ctx.GetRawContext(),                        /* context */
-      embed_dim,                                  /* embed_dim */
-      idx_len,                                    /* idx_len */
-      emb_layer_num,                              /* emb_layer_num */
-      param.padding_idx,                          /* padding_idx */
-      &arg_tables_[0],                            /* tables */
-      &arg_ids_[0],                               /* indices */
-      static_cast<int*>(table_lens_guard_.get()), /* table_lens */
-      nullptr,                                    /* scale_after_emb */
-      nullptr,                                    /* scale_after_ewadd */
+      ctx.GetRawContext(),                         /* context */
+      embed_dim,                                   /* embed_dim */
+      idx_len,                                     /* idx_len */
+      emb_layer_num,                               /* emb_layer_num */
+      param.padding_idx,                           /* padding_idx */
+      &arg_tables_[0],                             /* tables */
+      &arg_ids_[0],                                /* indices */
+      static_cast<int*>(table_lens_guard_->addr_), /* table_lens */
+      nullptr,                                     /* scale_after_emb */
+      nullptr,                                     /* scale_after_ewadd */
       param.Out->mutable_data<float>(TARGET(kXPU)) /* top */);
   CHECK_EQ(r, 0);
 }

lite/kernels/xpu/__xpu__embedding_with_eltwise_add_compute.h

@@ -14,10 +14,9 @@
 
 #pragma once
 
-#include <memory>
 #include <vector>
+#include "lite/backends/xpu/target_wrapper.h"  // XPUScratchPadGuard
 #include "lite/core/kernel.h"
-#include "lite/kernels/xpu/utils.h"  // XPUFreeDeleter
 
 namespace paddle {
 namespace lite {
@@ -36,7 +35,7 @@ class XPUEmbeddingWithEltwiseAddCompute
  private:
   std::vector<const int64_t*> arg_ids_;
   std::vector<const float*> arg_tables_;
-  std::unique_ptr<void, XPUFreeDeleter> table_lens_guard_;
+  XPUScratchPadGuard table_lens_guard_;
   std::vector<int> table_lens_cpu_;
 };
 

lite/kernels/xpu/__xpu__multi_encoder_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include <vector>
 #include "lite/backends/xpu/xpu_header_sitter.h"
 #include "lite/core/kernel.h"

lite/kernels/xpu/__xpu__resnet50_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include <vector>
 #include "lite/backends/xpu/xpu_header_sitter.h"
 #include "lite/core/kernel.h"

lite/kernels/xpu/__xpu__search_attention_compute.cc

@@ -22,16 +22,19 @@ namespace kernels {
 namespace xpu {
 
 void XPUMmdnnSearchAttentionCompute::PrepareForRun() {
-  offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  pad_begin_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  w_max_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(8 * sizeof(float));
+  offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  pad_begin_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  w_max_xpu_guard_ =
+      TargetWrapperXPU::MallocScratchPad(8 * sizeof(float), false /* use_l3 */);
   buffer_at_l3_guard_ = TargetWrapperXPU::MallocScratchPad(
       5 * L3_SLOT_SIZE * sizeof(float), false /* use_l3 */);
   buffer_at_gm_guard_ = TargetWrapperXPU::MallocScratchPad(
       5 * GM_SLOT_SIZE * sizeof(float), false /* use_l3 */);
 
-  offset_cpu.reset(new int[64]);
-  pad_begin_cpu.reset(new int[64]);
+  offset_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  pad_begin_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 void XPUMmdnnSearchAttentionCompute::Run() {
@@ -72,18 +75,18 @@ void XPUMmdnnSearchAttentionCompute::Run() {
   }
   offset_cpu[batch] = offset[batch];
 
-  xpu_memcpy(offset_xpu_guard_->addr_,
-             offset_cpu.get(),
-             offset.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(pad_begin_xpu_guard_->addr_,
-             pad_begin_cpu.get(),
-             batch * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(w_max_xpu_guard_->addr_,
-             maxs_cpu,
-             8 * sizeof(float),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(offset_xpu_guard_->addr_,
+                      offset_cpu.get(),
+                      offset.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(pad_begin_xpu_guard_->addr_,
+                      pad_begin_cpu.get(),
+                      batch * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(w_max_xpu_guard_->addr_,
+                      maxs_cpu,
+                      8 * sizeof(float),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int* offset_xpu = reinterpret_cast<int*>(offset_xpu_guard_->addr_);
   int* pad_begin_xpu = reinterpret_cast<int*>(pad_begin_xpu_guard_->addr_);
@@ -115,90 +118,99 @@ void XPUMmdnnSearchAttentionCompute::Run() {
   }
 
   const auto* bottom_data = X->data<float>();
-  xdnn::search_sequence_pad_depad(ctx.GetRawContext(),
-                                  const_cast<float*>(bottom_data),
-                                  group_padding_output,
-                                  offset_xpu,
-                                  max_seq,
-                                  batch,
-                                  dim1,
-                                  0);  // is_depad = 0
+  int r = 0;
+  r = xdnn::search_sequence_pad_depad(ctx.GetRawContext(),
+                                      const_cast<float*>(bottom_data),
+                                      group_padding_output,
+                                      offset_xpu,
+                                      max_seq,
+                                      batch,
+                                      dim1,
+                                      0);  // is_depad = 0
+  CHECK_EQ(r, 0);
   // do-findmax
-  xdnn::findmax<float>(ctx.GetRawContext(),
-                       group_padding_output,
-                       batch * max_seq * dim1,
-                       maxs_xpu);
-  xdnn::gemm_int16_maxptr<float, int16_t, float>(
-      ctx.GetRawContext(),
-      false,
-      true,  // trans_a, trans_b
-      batch * max_seq,
-      dim1,
-      dim1,  // m, n, k
-      1.0f,
-      group_padding_output,
-      dim1,  // alpha, data_a, lda
-      w_data,
-      dim1,
-      0.0f,  // data_b, ldb, beta
-      seq_fc_output,
-      dim1,
-      b_data,  // data_c, ldc, bias
-      xdnn::Activation_t::LINEAR,
-      maxs_xpu,
-      maxs_xpu + 4,
-      nullptr);  // max_a, max_b, max_c
-  xdnn::search_aligned_mat_mul(ctx.GetRawContext(),
-                               0,
-                               1,
-                               batch,
-                               max_seq,
-                               max_seq,
-                               dim1,
-                               alpha0,
-                               group_padding_output,
-                               dim1,
-                               seq_fc_output,
-                               dim1,
-                               batchgemm0_output,
-                               max_seq);
-  xdnn::search_pad_mask(ctx.GetRawContext(),
-                        batchgemm0_output,
-                        attention_output,
-                        pad_begin_xpu,
-                        batch,
-                        max_seq,
-                        max_seq,
-                        batch,
-                        mask);
-  xdnn::softmax2d_forward(ctx.GetRawContext(),
-                          attention_output,
-                          seq_softmax_output,
-                          batch * max_seq,
-                          max_seq,
-                          true);
-  xdnn::search_aligned_mat_mul(ctx.GetRawContext(),
-                               0,
-                               0,
-                               batch,
-                               max_seq,
-                               dim1,
-                               max_seq,
-                               alpha1,
-                               seq_softmax_output,
-                               max_seq,
-                               group_padding_output,
-                               dim1,
-                               batchgemm1_output,
-                               dim1);
-  xdnn::search_sequence_pad_depad(ctx.GetRawContext(),
-                                  top_data,
-                                  batchgemm1_output,
-                                  offset_xpu,
-                                  max_seq,
-                                  batch,
-                                  dim1,
-                                  1);  // is_depad = 1
+  r = xdnn::findmax<float>(ctx.GetRawContext(),
+                           group_padding_output,
+                           batch * max_seq * dim1,
+                           maxs_xpu);
+  CHECK_EQ(r, 0);
+  r = xdnn::gemm_int16_maxptr<float, int16_t, float>(
+      ctx.GetRawContext(),        /* ctx */
+      false,                      /* trans_a */
+      true,                       /* trans_b */
+      batch * max_seq,            /* m */
+      dim1,                       /* n */
+      dim1,                       /* k */
+      1.0f,                       /* alpha */
+      group_padding_output,       /* data_a */
+      dim1,                       /* lda */
+      w_data,                     /* data_b */
+      dim1,                       /* ldb */
+      0.0f,                       /* beta */
+      seq_fc_output,              /* data_c */
+      dim1,                       /* ldc */
+      b_data,                     /* bias */
+      xdnn::Activation_t::LINEAR, /* act */
+      maxs_xpu,                   /* max_a */
+      maxs_xpu + 4,               /* max_b */
+      nullptr /* max_c */);
+  CHECK_EQ(r, 0);
+  r = xdnn::search_aligned_mat_mul(ctx.GetRawContext(),
+                                   0,
+                                   1,
+                                   batch,
+                                   max_seq,
+                                   max_seq,
+                                   dim1,
+                                   alpha0,
+                                   group_padding_output,
+                                   dim1,
+                                   seq_fc_output,
+                                   dim1,
+                                   batchgemm0_output,
+                                   max_seq);
+  CHECK_EQ(r, 0);
+  r = xdnn::search_pad_mask(ctx.GetRawContext(),
+                            batchgemm0_output,
+                            attention_output,
+                            pad_begin_xpu,
+                            batch,
+                            max_seq,
+                            max_seq,
+                            batch,
+                            mask);
+  CHECK_EQ(r, 0);
+  r = xdnn::softmax2d_forward(ctx.GetRawContext(),
+                              attention_output,
+                              seq_softmax_output,
+                              batch * max_seq,
+                              max_seq,
+                              true);
+  CHECK_EQ(r, 0);
+  r = xdnn::search_aligned_mat_mul(ctx.GetRawContext(),
+                                   0,
+                                   0,
+                                   batch,
+                                   max_seq,
+                                   dim1,
+                                   max_seq,
+                                   alpha1,
+                                   seq_softmax_output,
+                                   max_seq,
+                                   group_padding_output,
+                                   dim1,
+                                   batchgemm1_output,
+                                   dim1);
+  CHECK_EQ(r, 0);
+  r = xdnn::search_sequence_pad_depad(ctx.GetRawContext(),
+                                      top_data,
+                                      batchgemm1_output,
+                                      offset_xpu,
+                                      max_seq,
+                                      batch,
+                                      dim1,
+                                      1);  // is_depad = 1
+  CHECK_EQ(r, 0);
 }
 
 }  // namespace xpu

lite/kernels/xpu/activation_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/batch_norm_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/cast_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/conv_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/dropout_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/elementwise_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/layer_norm_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/lookup_table_compute.cc

@@ -29,12 +29,13 @@ void LookupTableCompute::Run() {
   int embed_dim = param.W->dims()[1];
 
   int r = xdnn::embedding<float, int64_t>(
-      ctx.GetRawContext(),        /* context */
-      num,                        /* num */
-      param.Ids->data<int64_t>(), /* indices */
-      embed_dim,                  /* embed_dim */
-      param.W->data<float>(),     /* table */
-      param.Out->mutable_data<float>(TARGET(kXPU)) /* top */);
+      ctx.GetRawContext(),                          /* context */
+      num,                                          /* num */
+      param.Ids->data<int64_t>(),                   /* indices */
+      embed_dim,                                    /* embed_dim */
+      param.W->data<float>(),                       /* table */
+      param.Out->mutable_data<float>(TARGET(kXPU)), /* top */
+      param.padding_idx /* padding_idx */);
   CHECK_EQ(r, 0);
 }
 

lite/kernels/xpu/lookup_table_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/match_matrix_tensor_compute.cc

@@ -23,12 +23,15 @@ namespace kernels {
 namespace xpu {
 
 void MatchMatrixTensorCompute::PrepareForRun() {
-  wx_max_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  offset_l_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  offset_r_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-
-  offset_l_cpu.reset(new int[64]);
-  offset_r_cpu.reset(new int[64]);
+  wx_max_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  offset_l_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  offset_r_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+
+  offset_l_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  offset_r_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 void MatchMatrixTensorCompute::Run() {
@@ -76,25 +79,25 @@ void MatchMatrixTensorCompute::Run() {
   int* offset_r_xpu = reinterpret_cast<int*>(offset_r_xpu_guard_->addr_);
 
   int r = xdnn::gemm_int16_tmp_api<float, int16_t, float>(
-      ctx.GetRawContext(), /* ctx */
-      false,
-      false, /* trans_a, trans_b */
-      x->dims()[0],
-      dim_t * dim_in,
-      dim_in, /* m, n, k */
-      1.0f,
-      bottom_l_data,
-      dim_in, /* alpha, data_a, lda */
-      w_data,
-      dim_t * dim_in,
-      0.0f, /* data_b, ldb, beta */
-      bottom_l_trans_data,
-      dim_t * dim_in, /* data_c, ldc */
-      nullptr,        /* bias */
-      xdnn::Activation_t::LINEAR,
-      0.0f,
-      w_max,
-      wx_max /* max_a, max_b, max_c */);
+      ctx.GetRawContext(),        /* ctx */
+      false,                      /* trans_a */
+      false,                      /* trans_b */
+      x->dims()[0],               /* m */
+      dim_t * dim_in,             /* n */
+      dim_in,                     /* k */
+      1.0f,                       /* alpha */
+      bottom_l_data,              /* data_a */
+      dim_in,                     /* lda */
+      w_data,                     /* data_b */
+      dim_t * dim_in,             /* ldb */
+      0.0f,                       /* beta */
+      bottom_l_trans_data,        /* data_c */
+      dim_t * dim_in,             /* ldc */
+      nullptr,                    /* bias */
+      xdnn::Activation_t::LINEAR, /* act */
+      0.0f,                       /* max_a */
+      w_max,                      /* max_b */
+      wx_max /* max_c */);
   CHECK_EQ(r, 0);
 
   int max_width = 0;
@@ -110,14 +113,14 @@ void MatchMatrixTensorCompute::Run() {
       max_width = offset_r_cpu[i] - offset_r_cpu[i - 1];
     }
   }
-  xpu_memcpy(offset_l_xpu,
-             offset_l_cpu.get(),
-             offset_l.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(offset_r_xpu,
-             offset_r_cpu.get(),
-             offset_r.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(offset_l_xpu,
+                      offset_l_cpu.get(),
+                      offset_l.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(offset_r_xpu,
+                      offset_r_cpu.get(),
+                      offset_r.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   r = xdnn::match_matrix_tensor(ctx.GetRawContext(),
                                 batch_size,

lite/kernels/xpu/matmul_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/mul_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/pool_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/scale_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/search_fc_compute.cc

@@ -23,7 +23,8 @@ namespace kernels {
 namespace xpu {
 
 void SearchFcCompute::PrepareForRun() {
-  maxs_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(float));
+  maxs_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(float), false /* use_l3 */);
 }
 
 void SearchFcCompute::Run() {
@@ -59,34 +60,34 @@ void SearchFcCompute::Run() {
 
   float* maxs_xpu = reinterpret_cast<float*>(maxs_xpu_guard_->addr_);
   float maxs_cpu[8] = {0.0f, 0.0f, 0.0f, 0.0f, w_max, 0.0f, 0.0f, 0.0f};
-  xpu_memcpy(maxs_xpu,
-             &maxs_cpu[0],
-             8 * sizeof(float),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(maxs_xpu,
+                      &maxs_cpu[0],
+                      8 * sizeof(float),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::findmax<float>(
       ctx.GetRawContext(), bottom_data, batch * _in, maxs_xpu);
   CHECK_EQ(r, 0);
   r = xdnn::gemm_int16_maxptr<float, int16_t, float>(
       ctx.GetRawContext(), /* ctx */
-      false,
-      true, /*trans_a, trans_b*/
-      batch,
-      _out,
-      _in, /*m, n, k*/
-      1.0f,
-      bottom_data,
-      _in, /*alpha, data_a, lda*/
-      weights,
-      _in,
-      0.0f, /*data_b, ldb, beta*/
-      top_data,
-      _out,
-      bias_data, /* data_c, ldc, bias*/
-      act,
-      maxs_xpu,
-      maxs_xpu + 4,
-      nullptr /*act, max_a, max_b, max_c*/);
+      false,               /* trans_a */
+      true,                /* trans_b */
+      batch,               /* m */
+      _out,                /* n */
+      _in,                 /* k */
+      1.0f,                /* alpha */
+      bottom_data,         /* data_a */
+      _in,                 /* lda */
+      weights,             /* data_b */
+      _in,                 /* ldb */
+      0.0f,                /* beta */
+      top_data,            /* data_c */
+      _out,                /* ldc */
+      bias_data,           /* bias */
+      act,                 /* act */
+      maxs_xpu,            /* max_a */
+      maxs_xpu + 4,        /* max_b */
+      nullptr /* max_c */);
   CHECK_EQ(r, 0);
 }
 

lite/kernels/xpu/search_grnn_compute.cc

@@ -24,13 +24,16 @@ namespace kernels {
 namespace xpu {
 
 void SearchGrnnCompute::PrepareForRun() {
-  offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  new_offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(256 * sizeof(int));
-  maxs_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(16 * sizeof(float));
+  offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  new_offset_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SEQ_LEN * sizeof(int), false /* use_l3 */);
+  maxs_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(16 * sizeof(float),
+                                                       false /* use_l3 */);
 
-  idx_sorted_by_width_data_cpu.reset(new int[64]);
-  offset_cpu.reset(new int[64]);
-  new_offset_cpu.reset(new int[256]);
+  idx_sorted_by_width_data_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  offset_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  new_offset_cpu.reset(new int[XPU_MAX_LOD_SEQ_LEN]);
 }
 
 void SearchGrnnCompute::prepare_layout(const operators::SearchGrnnParam& param,
@@ -96,10 +99,10 @@ void SearchGrnnCompute::prepare_layout(const operators::SearchGrnnParam& param,
     layout_input->Resize({dim0, dim1});
   }
 
-  xpu_memcpy(idx_sorted_by_width->mutable_data<int>(TARGET(kXPU)),
-             idx_sorted_by_width_data_cpu.get(),
-             idx_sorted_by_width->numel() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(idx_sorted_by_width->mutable_data<int>(TARGET(kXPU)),
+                      idx_sorted_by_width_data_cpu.get(),
+                      idx_sorted_by_width->numel() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 }
 
 void SearchGrnnCompute::Run() {
@@ -156,14 +159,14 @@ void SearchGrnnCompute::Run() {
   for (size_t i = 0; i < new_offset.size(); ++i) {
     new_offset_cpu[i] = new_offset[i];
   }
-  xpu_memcpy(offset_xpu,
-             offset_cpu.get(),
-             offset.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(new_offset_xpu,
-             new_offset_cpu.get(),
-             new_offset.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(offset_xpu,
+                      offset_cpu.get(),
+                      offset.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(new_offset_xpu,
+                      new_offset_cpu.get(),
+                      new_offset.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::search_seq2batch(ctx.GetRawContext(),
                                  batch,
@@ -200,10 +203,10 @@ void SearchGrnnCompute::Run() {
                         0.0f,
                         0.0f,
                         0.0f};
-  xpu_memcpy(maxs_xpu,
-             maxs_cpu,
-             16 * sizeof(float),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(maxs_xpu,
+                      maxs_cpu,
+                      16 * sizeof(float),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
   r = xdnn::findmax<float>(
       ctx.GetRawContext(), new_emb, cap_l * cap_e, maxs_xpu);
   CHECK_EQ(r, 0);

lite/kernels/xpu/sequence_arithmetic_compute.cc

@@ -37,44 +37,54 @@ void SequenceArithmeticCompute::Run() {
   const auto* bottom_data1 = bottom1->data<float>();
   auto* top_data = top->mutable_data<float>(TARGET(kXPU));
 
+  int r = 0;
   switch (op_type) {
     case 1:  // addition: top[0] = bottom[0] + bottom[1]
       if (len1 > len2) {
-        xdnn::elementwise_add(
+        r = xdnn::elementwise_add(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len2);
-        xdnn::memcpy_device(ctx.GetRawContext(),
-                            &top_data[len2],
-                            &bottom_data0[len2],
-                            (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
+        r = xdnn::memcpy_device(ctx.GetRawContext(),
+                                &top_data[len2],
+                                &bottom_data0[len2],
+                                (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
       } else {
-        xdnn::elementwise_add(
+        r = xdnn::elementwise_add(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len1);
+        CHECK_EQ(r, 0);
       }
       break;
     case 2:  // substraction: top[0] = bottom[0] - bottom[1]
       if (len1 > len2) {
-        xdnn::elementwise_sub(
+        r = xdnn::elementwise_sub(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len2);
-        xdnn::memcpy_device(ctx.GetRawContext(),
-                            &top_data[len2],
-                            &bottom_data0[len2],
-                            (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
+        r = xdnn::memcpy_device(ctx.GetRawContext(),
+                                &top_data[len2],
+                                &bottom_data0[len2],
+                                (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
       } else {
-        xdnn::elementwise_sub(
+        r = xdnn::elementwise_sub(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len1);
+        CHECK_EQ(r, 0);
       }
       break;
     case 3:  // multiplication: top[0] = bottom[0] * bottom[1]
       if (len1 > len2) {
-        xdnn::elementwise_mul(
+        r = xdnn::elementwise_mul(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len2);
-        xdnn::memcpy_device(ctx.GetRawContext(),
-                            &top_data[len2],
-                            &bottom_data0[len2],
-                            (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
+        r = xdnn::memcpy_device(ctx.GetRawContext(),
+                                &top_data[len2],
+                                &bottom_data0[len2],
+                                (len1 - len2) * sizeof(float));
+        CHECK_EQ(r, 0);
       } else {
-        xdnn::elementwise_mul(
+        r = xdnn::elementwise_mul(
             ctx.GetRawContext(), bottom_data0, bottom_data1, top_data, len1);
+        CHECK_EQ(r, 0);
       }
       break;
     default:

lite/kernels/xpu/sequence_concat_compute.cc

@@ -23,11 +23,13 @@ namespace kernels {
 namespace xpu {
 
 void SequenceConcatCompute::PrepareForRun() {
-  lod0_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  lod1_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
+  lod0_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  lod1_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
 
-  lod0_cpu.reset(new int[64]);
-  lod1_cpu.reset(new int[64]);
+  lod0_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  lod1_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 template <typename T>
@@ -106,14 +108,14 @@ void SequenceConcatCompute::Run() {
   for (int i = 0; i < lod1.size(); ++i) {
     lod1_cpu[i] = lod1[i];
   }
-  xpu_memcpy(lod0_xpu,
-             lod0_cpu.get(),
-             lod0.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(lod1_xpu,
-             lod1_cpu.get(),
-             lod1.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(lod0_xpu,
+                      lod0_cpu.get(),
+                      lod0.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(lod1_xpu,
+                      lod1_cpu.get(),
+                      lod1.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::sequence_concat(ctx.GetRawContext(),
                                 xs[0]->data<float>(),

lite/kernels/xpu/sequence_pool_compute.cc

@@ -23,8 +23,9 @@ namespace kernels {
 namespace xpu {
 
 void XPUSequencePoolCompute::PrepareForRun() {
-  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  lod_cpu.reset(new int[64]);
+  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  lod_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 void XPUSequencePoolCompute::Run() {
@@ -55,10 +56,10 @@ void XPUSequencePoolCompute::Run() {
     lod_cpu[i] = in_lod[i];
   }
   int* lod_xpu = reinterpret_cast<int*>(lod_xpu_guard_->addr_);
-  xpu_memcpy(lod_xpu,
-             lod_cpu.get(),
-             in_lod.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(lod_xpu,
+                      lod_cpu.get(),
+                      in_lod.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r =
       xdnn::sequence_pooling_forward(ctx.GetRawContext(),

lite/kernels/xpu/sequence_reverse_compute.cc

@@ -23,8 +23,9 @@ namespace xpu {
 
 template <typename T, PrecisionType PType>
 void SequenceReverseCompute<T, PType>::PrepareForRun() {
-  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  lod_cpu.reset(new int[64]);
+  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  lod_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 template <typename T, PrecisionType PType>
@@ -58,10 +59,10 @@ void SequenceReverseCompute<T, PType>::Run() {
     lod_cpu[i] = lod[i];
   }
   int* lod_xpu = reinterpret_cast<int*>(lod_xpu_guard_->addr_);
-  xpu_memcpy(lod_xpu,
-             lod_cpu.get(),
-             lod.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(lod_xpu,
+                      lod_cpu.get(),
+                      lod.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::sequence_reverse(ctx.GetRawContext(),
                                  batch_size,

lite/kernels/xpu/sequence_topk_avg_pooling_compute.cc

@@ -23,10 +23,11 @@ namespace kernels {
 namespace xpu {
 
 void SequenceTopkAvgPoolingCompute::PrepareForRun() {
-  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(256 * sizeof(int));
-  in_lod_cpu.reset(new int[64]);
-  row_lod_cpu.reset(new int[64]);
-  col_lod_cpu.reset(new int[64]);
+  lod_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      4 * XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  in_lod_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  row_lod_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  col_lod_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 void SequenceTopkAvgPoolingCompute::Run() {
@@ -81,22 +82,22 @@ void SequenceTopkAvgPoolingCompute::Run() {
   for (int i = 0; i < col_lod.size(); ++i) {
     col_lod_cpu[i] = col_lod[i];
   }
-  xpu_memcpy(in_lod_xpu,
-             in_lod_cpu.get(),
-             in_lod.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(row_lod_xpu,
-             row_lod_cpu.get(),
-             row_lod.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(col_lod_xpu,
-             col_lod_cpu.get(),
-             col_lod.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(topks_xpu,
-             topks.data(),
-             topks.size() * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(in_lod_xpu,
+                      in_lod_cpu.get(),
+                      in_lod.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(row_lod_xpu,
+                      row_lod_cpu.get(),
+                      row_lod.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(col_lod_xpu,
+                      col_lod_cpu.get(),
+                      col_lod.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(topks_xpu,
+                      topks.data(),
+                      topks.size() * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::sequence_topk_avg_pooling(ctx.GetRawContext(),
                                           in_data,

lite/kernels/xpu/softmax_compute.h

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #pragma once
+
 #include "lite/core/kernel.h"
 
 namespace paddle {

lite/kernels/xpu/stack_compute.cc

@@ -25,9 +25,8 @@ void StackCompute::PrepareForRun() {
   auto& param = this->Param<param_t>();
 
   int n = param.X.size();
-  void* x_ptr = nullptr;
-  xpu_malloc(&x_ptr, n * 8 /* sizeof(__global__ float*) */);
-  x_ptr_guard_.reset(x_ptr);
+  x_ptr_guard_ = TargetWrapperXPU::MallocScratchPad(
+      n * 8 /* sizeof(__global__ float*) */, false /* use_l3 */);
   x_ptr_cpu_.reserve(n);
 }
 
@@ -47,14 +46,15 @@ void StackCompute::Run() {
   for (int i = 0; i < n; ++i) {
     x_ptr_cpu_[i] = param.X[i]->data<float>();
   }
-  xpu_memcpy(x_ptr_guard_.get(), &x_ptr_cpu_[0], n * 8, XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(
+      x_ptr_guard_->addr_, &x_ptr_cpu_[0], n * 8, XPU_HOST_TO_DEVICE));
 
   int r = xdnn::stack_forward(
       ctx.GetRawContext(), /* context */
       height,              /* height */
       width,               /* width */
       n,                   /* n */
-      x_ptr_guard_.get(),  /* x_ptr */
+      x_ptr_guard_->addr_, /* x_ptr */
       param.Out->mutable_data<float>(TARGET(kXPU)) /* out */);
   CHECK_EQ(r, 0);
 }

lite/kernels/xpu/stack_compute.h

@@ -14,10 +14,9 @@
 
 #pragma once
 
-#include <memory>
 #include <vector>
+#include "lite/backends/xpu/target_wrapper.h"  // XPUScratchPadGuard
 #include "lite/core/kernel.h"
-#include "lite/kernels/xpu/utils.h"  // XPUFreeDeleter
 
 namespace paddle {
 namespace lite {
@@ -35,7 +34,7 @@ class StackCompute : public KernelLite<TARGET(kXPU), PRECISION(kFloat)> {
   virtual ~StackCompute() = default;
 
  private:
-  std::unique_ptr<void, XPUFreeDeleter> x_ptr_guard_;
+  XPUScratchPadGuard x_ptr_guard_;
   std::vector<const float*> x_ptr_cpu_;
 };
 

lite/kernels/xpu/utils.h

-// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
-//
-// 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.
-
-#pragma once
-
-#include "lite/backends/xpu/xpu_header_sitter.h"
-
-namespace paddle {
-namespace lite {
-namespace kernels {
-namespace xpu {
-
-struct XPUFreeDeleter {
-  void operator()(void* p) const { xpu_free(p); }
-};
-
-}  // namespace xpu
-}  // namespace kernels
-}  // namespace lite
-}  // namespace paddle

lite/kernels/xpu/var_conv_2d_compute.cc

@@ -23,10 +23,12 @@ namespace kernels {
 namespace xpu {
 
 void VarConv2DCompute::PrepareForRun() {
-  offset_x_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  offset_y_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(64 * sizeof(int));
-  offset_x_cpu.reset(new int[64]);
-  offset_y_cpu.reset(new int[64]);
+  offset_x_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  offset_y_xpu_guard_ = TargetWrapperXPU::MallocScratchPad(
+      XPU_MAX_LOD_SIZE * sizeof(int), false /* use_l3 */);
+  offset_x_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
+  offset_y_cpu.reset(new int[XPU_MAX_LOD_SIZE]);
 }
 
 void VarConv2DCompute::Run() {
@@ -94,14 +96,14 @@ void VarConv2DCompute::Run() {
     offset_x_cpu[i] = offset_x[i];
     offset_y_cpu[i] = offset_y[i];
   }
-  xpu_memcpy(offset_x_xpu,
-             offset_x_cpu.get(),
-             (batch + 1) * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
-  xpu_memcpy(offset_y_xpu,
-             offset_y_cpu.get(),
-             (batch + 1) * sizeof(int),
-             XPUMemcpyKind::XPU_HOST_TO_DEVICE);
+  XPU_CALL(xpu_memcpy(offset_x_xpu,
+                      offset_x_cpu.get(),
+                      (batch + 1) * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
+  XPU_CALL(xpu_memcpy(offset_y_xpu,
+                      offset_y_cpu.get(),
+                      (batch + 1) * sizeof(int),
+                      XPUMemcpyKind::XPU_HOST_TO_DEVICE));
 
   int r = xdnn::search_varconv<float, int16_t>(ctx.GetRawContext(),
                                                batch,

lite/operators/__xpu__mmdnn_op.cc

@@ -88,6 +88,78 @@ bool XPUMmdnnBidEmbGrnnAttOp::AttachImpl(const cpp::OpDesc& op_desc,
   return true;
 }
 
+bool XPUMmdnnBidEmbGrnnAttOp2::CheckShape() const { return true; }
+
+bool XPUMmdnnBidEmbGrnnAttOp2::InferShapeImpl() const {
+  auto& id_dims = param_.id0->dims();
+  auto& id_lod = param_.id0->lod()[0];
+  auto& emb_tbl_dims = param_.emb_tbl->dims();
+  auto& grnn_wh_dims = param_.grnn_rv_wh->dims();
+
+  param_.emb0_out->Resize({id_dims[0], emb_tbl_dims[1]});
+  param_.emb0_out->set_lod({id_lod});
+  param_.grnn_fw_pool_out->Resize(
+      {(int64_t)id_lod.size() - 1, grnn_wh_dims[2]});
+  param_.grnn_rv_pool_out->Resize(
+      {(int64_t)id_lod.size() - 1, grnn_wh_dims[2]});
+  param_.att_pool_out->Resize(
+      {(int64_t)id_lod.size() - 1, 2 * grnn_wh_dims[2]});
+  param_.concat_3in1_out->Resize({id_dims[0], 3 * grnn_wh_dims[2]});
+  param_.concat_3in1_out->set_lod({id_lod});
+  param_.emb_fw_out->Resize({id_dims[0], emb_tbl_dims[1]});
+  param_.emb_fw_out->set_lod({id_lod});
+  return true;
+}
+
+bool XPUMmdnnBidEmbGrnnAttOp2::AttachImpl(const cpp::OpDesc& op_desc,
+                                          lite::Scope* scope) {
+  param_.id0 =
+      scope->FindVar(op_desc.Input("id0").front())->GetMutable<lite::Tensor>();
+  param_.id1 =
+      scope->FindVar(op_desc.Input("id1").front())->GetMutable<lite::Tensor>();
+  param_.emb_tbl = scope->FindVar(op_desc.Input("emb_tbl").front())
+                       ->GetMutable<lite::Tensor>();
+  param_.grnn_fw_wh = scope->FindVar(op_desc.Input("grnn_fw_wh").front())
+                          ->GetMutable<lite::Tensor>();
+  param_.grnn_fw_wi = scope->FindVar(op_desc.Input("grnn_fw_wi").front())
+                          ->GetMutable<lite::Tensor>();
+  param_.grnn_rv_wh = scope->FindVar(op_desc.Input("grnn_rv_wh").front())
+                          ->GetMutable<lite::Tensor>();
+  param_.grnn_rv_wi = scope->FindVar(op_desc.Input("grnn_rv_wi").front())
+                          ->GetMutable<lite::Tensor>();
+  param_.att_fc_w = scope->FindVar(op_desc.Input("att_fc_w").front())
+                        ->GetMutable<lite::Tensor>();
+  param_.att_fc_b = scope->FindVar(op_desc.Input("att_fc_b").front())
+                        ->GetMutable<lite::Tensor>();
+
+  param_.emb0_out = scope->FindVar(op_desc.Output("emb0_out").front())
+                        ->GetMutable<lite::Tensor>();
+  param_.grnn_fw_pool_out =
+      scope->FindVar(op_desc.Output("grnn_fw_pool_out").front())
+          ->GetMutable<lite::Tensor>();
+  param_.grnn_rv_pool_out =
+      scope->FindVar(op_desc.Output("grnn_rv_pool_out").front())
+          ->GetMutable<lite::Tensor>();
+  param_.att_pool_out = scope->FindVar(op_desc.Output("att_pool_out").front())
+                            ->GetMutable<lite::Tensor>();
+  param_.concat_3in1_out =
+      scope->FindVar(op_desc.Output("concat_3in1_out").front())
+          ->GetMutable<lite::Tensor>();
+  param_.emb_fw_out = scope->FindVar(op_desc.Output("emb_fw_out").front())
+                          ->GetMutable<lite::Tensor>();
+
+  param_.grnn_fw_wh_maxs =
+      op_desc.GetAttr<std::vector<float>>("grnn_fw_wh_maxs");
+  param_.grnn_fw_wi_maxs =
+      op_desc.GetAttr<std::vector<float>>("grnn_fw_wi_maxs");
+  param_.grnn_rv_wh_maxs =
+      op_desc.GetAttr<std::vector<float>>("grnn_rv_wh_maxs");
+  param_.grnn_rv_wi_maxs =
+      op_desc.GetAttr<std::vector<float>>("grnn_rv_wi_maxs");
+  param_.att_fc_w_max = op_desc.GetAttr<float>("att_fc_w_max");
+  return true;
+}
+
 bool XPUMmdnnBidEmbAttOp::CheckShape() const { return true; }
 
 bool XPUMmdnnBidEmbAttOp::InferShapeImpl() const {
@@ -157,6 +229,7 @@ bool XPUMmdnnMatchConvTopkOp::AttachImpl(const cpp::OpDesc& op_desc,
   param_.input_w_max = op_desc.GetAttr<float>("input_w_max");
   param_.conv_w_max = op_desc.GetAttr<float>("conv_w_max");
   param_.topks = op_desc.GetAttr<std::vector<int>>("topks");
+  param_.output_channel = op_desc.GetAttr<int>("output_channel");
   param_.channel_num = op_desc.GetAttr<int>("channel_num");
   param_.dim_t = op_desc.GetAttr<int>("dim_t");
   return true;
@@ -182,10 +255,10 @@ bool XPUMmdnnMergeAllOp::AttachImpl(const cpp::OpDesc& op_desc,
     auto t = scope->FindVar(name)->GetMutable<lite::Tensor>();
     param_.concat_7in1_x.push_back(t);
   }
-  param_.concat_2in1_x.clear();
-  for (auto& name : op_desc.Input("concat_2in1_x")) {
+  param_.concat_topk_x.clear();
+  for (auto& name : op_desc.Input("concat_topk_x")) {
     auto t = scope->FindVar(name)->GetMutable<lite::Tensor>();
-    param_.concat_2in1_x.push_back(t);
+    param_.concat_topk_x.push_back(t);
   }
   param_.grnn_fw_wh = scope->FindVar(op_desc.Input("grnn_fw_wh").front())
                           ->GetMutable<lite::Tensor>();
@@ -231,6 +304,8 @@ bool XPUMmdnnMergeAllOp::AttachImpl(const cpp::OpDesc& op_desc,
 
 REGISTER_LITE_OP(__xpu__mmdnn_bid_emb_grnn_att,
                  paddle::lite::operators::XPUMmdnnBidEmbGrnnAttOp);
+REGISTER_LITE_OP(__xpu__mmdnn_bid_emb_grnn_att2,
+                 paddle::lite::operators::XPUMmdnnBidEmbGrnnAttOp2);
 REGISTER_LITE_OP(__xpu__mmdnn_bid_emb_att,
                  paddle::lite::operators::XPUMmdnnBidEmbAttOp);
 REGISTER_LITE_OP(__xpu__mmdnn_match_conv_topk,

lite/operators/__xpu__mmdnn_op.h

@@ -41,6 +41,29 @@ class XPUMmdnnBidEmbGrnnAttOp : public OpLite {
   mutable XPUMmdnnBidEmbGrnnAttParam param_;
 };
 
+class XPUMmdnnBidEmbGrnnAttOp2 : public OpLite {
+ public:
+  XPUMmdnnBidEmbGrnnAttOp2() {}
+
+  explicit XPUMmdnnBidEmbGrnnAttOp2(const std::string &op_type)
+      : OpLite(op_type) {}
+
+  bool CheckShape() const override;
+
+  bool InferShapeImpl() const override;
+
+  bool AttachImpl(const cpp::OpDesc &opdesc, lite::Scope *scope) override;
+
+  void AttachKernel(KernelBase *kernel) override { kernel->SetParam(param_); }
+
+  std::string DebugString() const override {
+    return "XPUMmdnnBidEmbGrnnAttOp2";
+  }
+
+ private:
+  mutable XPUMmdnnBidEmbGrnnAttParam2 param_;
+};
+
 class XPUMmdnnBidEmbAttOp : public OpLite {
  public:
   XPUMmdnnBidEmbAttOp() {}

lite/operators/op_params.h

@@ -1627,11 +1627,36 @@ struct XPUMmdnnBidEmbGrnnAttParam : ParamBase {
   std::vector<float> grnn_rv_wi_maxs;
   float att_fc_w_max{0.0f};
 
-  lite::Tensor* grnn_fw_pool_out{};  // 1
-  lite::Tensor* grnn_rv_pool_out{};  // 2
-  lite::Tensor* att_pool_out{};      // 3
-  lite::Tensor* concat_3in1_out{};   // 4
-  lite::Tensor* emb_fw_out{};        // 5
+  lite::Tensor* grnn_fw_pool_out{};
+  lite::Tensor* grnn_rv_pool_out{};
+  lite::Tensor* att_pool_out{};
+  lite::Tensor* concat_3in1_out{};
+  lite::Tensor* emb_fw_out{};
+};
+
+struct XPUMmdnnBidEmbGrnnAttParam2 : ParamBase {
+  lite::Tensor* id0{};
+  lite::Tensor* id1{};
+  lite::Tensor* emb_tbl{};
+  lite::Tensor* grnn_fw_wh{};
+  lite::Tensor* grnn_fw_wi{};
+  lite::Tensor* grnn_rv_wh{};
+  lite::Tensor* grnn_rv_wi{};
+  lite::Tensor* att_fc_w{};
+  lite::Tensor* att_fc_b{};
+
+  std::vector<float> grnn_fw_wh_maxs;
+  std::vector<float> grnn_fw_wi_maxs;
+  std::vector<float> grnn_rv_wh_maxs;
+  std::vector<float> grnn_rv_wi_maxs;
+  float att_fc_w_max{0.0f};
+
+  lite::Tensor* emb0_out{};
+  lite::Tensor* grnn_fw_pool_out{};
+  lite::Tensor* grnn_rv_pool_out{};
+  lite::Tensor* att_pool_out{};
+  lite::Tensor* concat_3in1_out{};
+  lite::Tensor* emb_fw_out{};
 };
 
 struct XPUMmdnnBidEmbAttParam : ParamBase {
@@ -1643,8 +1668,8 @@ struct XPUMmdnnBidEmbAttParam : ParamBase {
 
   float att_fc_w_max{0.0f};
 
-  lite::Tensor* att_pool_out{};  // 1
-  lite::Tensor* emb_fw_out{};    // 2
+  lite::Tensor* att_pool_out{};
+  lite::Tensor* emb_fw_out{};
 };
 
 struct XPUMmdnnMatchConvTopkParam : ParamBase {
@@ -1656,6 +1681,7 @@ struct XPUMmdnnMatchConvTopkParam : ParamBase {
   float input_w_max{0.0f};
   float conv_w_max{0.0f};
   std::vector<int> topks;
+  int output_channel{0};
   int channel_num{0};
   int dim_t{0};
 
@@ -1664,7 +1690,7 @@ struct XPUMmdnnMatchConvTopkParam : ParamBase {
 
 struct XPUMmdnnMergeAllParam : ParamBase {
   std::vector<lite::Tensor*> concat_7in1_x;
-  std::vector<lite::Tensor*> concat_2in1_x;
+  std::vector<lite::Tensor*> concat_topk_x;
   lite::Tensor* grnn_fw_wh{};
   lite::Tensor* grnn_fw_wi{};
   lite::Tensor* grnn_rv_wh{};