[CINN] prefetch load scalar tensor for vectorize situation

Author: ZhangX-21

paddle/cinn/hlir/framework/pir/op_lowering_impl.cc

@@ -433,6 +433,11 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(
     } else {
       func = optim::Optimize(func, common::DefaultHostTarget(), false);
     }
+    auto pre_load_temp_buffers =
+        lang::GetPreLoadTempBufferAfterVectorize(func->body);
+    func->temp_bufs.insert(func->temp_bufs.end(),
+                           pre_load_temp_buffers.begin(),
+                           pre_load_temp_buffers.end());
     func->num_output_tensors = infer_shape_arg_tensor->size();
     lowered_funcs.push_back(std::move(func));
   }

paddle/cinn/lang/lower.cc

@@ -228,6 +228,22 @@ std::vector<ir::Buffer> GetTempBuffers(const std::vector<ir::Argument>& args,
   return temp_buffers;
 }
 
+std::vector<ir::Buffer> GetPreLoadTempBufferAfterVectorize(Expr body) {
+  std::unordered_set<std::string> buffer_names;
+  std::vector<ir::Buffer> temp_buffers;
+  ir::ir_utils::CollectIRNodesWithoutTensor(body, [&](const Expr* x) {
+    if (x->as_tensor() && x->as_tensor()->buffer.defined() &&
+        !buffer_names.count(x->as_tensor()->buffer->name) &&
+        utils::StartsWith(x->as_tensor()->buffer->name, "pre_load")) {
+      buffer_names.insert(x->as_tensor()->buffer->name);
+      temp_buffers.push_back(x->as_tensor()->buffer);
+      return true;
+    }
+    return false;
+  });
+  return std::move(temp_buffers);
+}
+
 std::set<ir::Tensor> CollectTempTensorsFromCtrlDepends(
     ast_gen_ius::TensorGroup* tensor_group,
     const std::vector<Tensor>& tensor_args) {

paddle/cinn/lang/lower.h

@@ -59,5 +59,7 @@ std::vector<ir::Buffer> GetTempBuffers(const std::vector<ir::Argument> &args,
 std::vector<ir::Buffer> GetTempBuffers(
     const std::vector<cinn::ir::Tensor> &tensor_args, Expr body);
 
+std::vector<ir::Buffer> GetPreLoadTempBufferAfterVectorize(Expr body);
+
 }  // namespace lang
 }  // namespace cinn