[XPU] Fix argmax, strided_slice when L3 cache and auto-tune is enabled. (#10121) (#10125)

Author: MuShangCC

lite/kernels/xpu/argmax_compute.cc

@@ -22,6 +22,15 @@ namespace lite {
 namespace kernels {
 namespace xpu {
 
+void ArgmaxCompute::PrepareForRun() {
+  auto& param = this->template Param<param_t>();
+  if (param.dtype == 2) {
+    auto out = param.Out;
+    out_int64_xpu_guard_ =
+        TargetWrapperXPU::MallocScratchPad(out->numel() * sizeof(int64_t));
+  }
+}
+
 void ArgmaxCompute::Run() {
   auto& param = this->template Param<param_t>();
   auto& ctx = this->ctx_->template As<XPUContext>();
@@ -45,19 +54,15 @@ void ArgmaxCompute::Run() {
     CHECK_EQ(r, 0);
   } else if (param.dtype == 2) {
     // int32
-    Tensor out_int64;
-    out_int64.Resize(out->dims());
+    int64_t* out_int64_data =
+        reinterpret_cast<int64_t*>(out_int64_xpu_guard_->addr_);
     int r = xdnn::argmax<float, int64_t>(
-        ctx.GetRawContext(),
-        x->data<float>(),
-        out_int64.mutable_data<int64_t>(TARGET(kXPU)),
-        x_dims,
-        axis);
+        ctx.GetRawContext(), x->data<float>(), out_int64_data, x_dims, axis);
     CHECK_EQ(r, 0);
     r = xdnn::cast_v2<int64_t, int>(ctx.GetRawContext(),
-                                    out_int64.data<int64_t>(),
+                                    out_int64_data,
                                     out->mutable_data<int>(TARGET(kXPU)),
-                                    out_int64.numel());
+                                    static_cast<int>(out->numel()));
     CHECK_EQ(r, 0);
   } else {
     LOG(FATAL) << "argmax unsupported param type for xpu: " << param.dtype;

lite/kernels/xpu/argmax_compute.h

@@ -25,9 +25,14 @@ class ArgmaxCompute : public KernelLite<TARGET(kXPU), PRECISION(kFloat)> {
  public:
   using param_t = operators::ArgmaxParam;
 
+  void PrepareForRun() override;
+
   virtual void Run();
 
   virtual ~ArgmaxCompute() = default;
+
+ private:
+  XPUScratchPadGuard out_int64_xpu_guard_;
 };
 
 }  // namespace xpu

lite/kernels/xpu/strided_slice_compute.cc

@@ -296,9 +296,9 @@ void StridedSliceCompute<T, PType>::Run() {
   auto* out_t = param.Out->template mutable_data<T>(TARGET(kXPU));
 
   if (need_reverse) {
-    lite::Tensor* tmp = new lite::Tensor();
-    tmp->Resize(out_dims);
-    auto* tmp_t = tmp->template mutable_data<T>(TARGET(kXPU));
+    XPUScratchPadGuard tmp_xpu_guard =
+        TargetWrapperXPU::MallocScratchPad(param.Out->numel() * sizeof(T));
+    auto tmp_t = reinterpret_cast<T*>(tmp_xpu_guard->addr_);
     int r = xdnn::strided_slice<T>(ctx.GetRawContext(),
                                    in_t,
                                    tmp_t,