reorder load sotre process (#61901)

phlrain · web-flow · commit 6e37e9b34232 · 2024-02-22T09:38:15.000+08:00
diff --git a/paddle/cinn/optim/transform_gpu_forloop.cc b/paddle/cinn/optim/transform_gpu_forloop.cc
@@ -319,6 +319,8 @@ class LocalAxisVisitor : public ir::IRMutator<> {
  private:
   void Visit(const ir::Store *op, Expr *expr) override {
     auto store = expr->As<ir::Store>();
+
+    ir::IRMutator<>::Visit(op, expr);
     if (!store->tensor.as_tensor_ref()->buffer.defined()) {
       return;
     }
@@ -332,11 +334,11 @@ class LocalAxisVisitor : public ir::IRMutator<> {
         indice = cinn::common::AutoSimplify(indice);
       }
     }
-    ir::IRMutator<>::Visit(op, expr);
   }
 
   void Visit(const ir::Load *op, Expr *expr) override {
     auto load = expr->As<ir::Load>();
+
     if (load->is_addr_scalar()) {
       return;
     }

Original file line number	Diff line number	Diff line change
`@@ -319,6 +319,8 @@ class LocalAxisVisitor : public ir::IRMutator<> {`
`319`	`319`	`private:`
`320`	`320`	`void Visit(const ir::Store op, Expr expr) override {`
`321`	`321`	`auto store = expr->As<ir::Store>();`
	`322`	`+`
	`323`	`+ ir::IRMutator<>::Visit(op, expr);`
`322`	`324`	`if (!store->tensor.as_tensor_ref()->buffer.defined()) {`
`323`	`325`	`return;`
`324`	`326`	`}`
`@@ -332,11 +334,11 @@ class LocalAxisVisitor : public ir::IRMutator<> {`
`332`	`334`	`indice = cinn::common::AutoSimplify(indice);`
`333`	`335`	`}`
`334`	`336`	`}`
`335`		`- ir::IRMutator<>::Visit(op, expr);`
`336`	`337`	`}`
`337`	`338`
`338`	`339`	`void Visit(const ir::Load op, Expr expr) override {`
`339`	`340`	`auto load = expr->As<ir::Load>();`
	`341`	`+`
`340`	`342`	`if (load->is_addr_scalar()) {`
`341`	`343`	`return;`
`342`	`344`	`}`