[CINN] prefetch load scalar tensor for vectorize situation (#71249)

ZhangX-21 · web-flow · commit 79bdacc3fb16 · 2025-03-03T23:16:09.000+08:00
diff --git a/paddle/cinn/hlir/framework/pir/op_lowering_impl.cc b/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));
   }
diff --git a/paddle/cinn/lang/lower.cc b/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) {
diff --git a/paddle/cinn/lang/lower.h b/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
diff --git a/paddle/cinn/optim/vectorize_for_trans.cc b/paddle/cinn/optim/vectorize_for_trans.cc

Original file line number	Diff line number	Diff line change
`@@ -433,6 +433,11 @@ std::vector<ir::LoweredFunc> OpLowererImpl::PostProcess(`
`433`	`433`	`} else {`
`434`	`434`	`func = optim::Optimize(func, common::DefaultHostTarget(), false);`
`435`	`435`	`}`
	`436`	`+ auto pre_load_temp_buffers =`
	`437`	`+ lang::GetPreLoadTempBufferAfterVectorize(func->body);`
	`438`	`+ func->temp_bufs.insert(func->temp_bufs.end(),`
	`439`	`+ pre_load_temp_buffers.begin(),`
	`440`	`+ pre_load_temp_buffers.end());`
`436`	`441`	`func->num_output_tensors = infer_shape_arg_tensor->size();`
`437`	`442`	`lowered_funcs.push_back(std::move(func));`
`438`	`443`	`}`