Fix matcher

zhang677 · zhang677 · commit 79dd27c54242 · 2022-07-01T16:47:39.000+08:00
diff --git a/src/index_notation/index_notation.cpp b/src/index_notation/index_notation.cpp
@@ -2055,13 +2055,15 @@ IndexStmt IndexStmt::wsaccel(TensorVar& ws, const std::vector<IndexVar>& accels,
     }
     set<IndexVar> TempVars;
     match(*this,
-          std::function<void(const WhereNode*)>([&](const WhereNode* where) {
+          std::function<void(const WhereNode*,Matcher*)>([&](const WhereNode* where,Matcher* ctx) {
         auto Temp = getResultAccesses(where->producer).first[0];
         if (Temp.getTensorVar() == ws) {
             for (auto i :getIndexVars()){
                 TempVars.insert(i);
             }
         }
+        ctx->match(where->producer);
+        ctx->match(where->consumer);
     }));
     for (auto i : accels) {
         if (TempVars.find(i) == TempVars.end()) {

Original file line number	Diff line number	Diff line change
`@@ -2055,13 +2055,15 @@ IndexStmt IndexStmt::wsaccel(TensorVar& ws, const std::vector<IndexVar>& accels,`
`2055`	`2055`	`}`
`2056`	`2056`	`set<IndexVar> TempVars;`
`2057`	`2057`	`match(*this,`
`2058`		`- std::function<void(const WhereNode)>([&](const WhereNode where) {`
	`2058`	`+ std::function<void(const WhereNode,Matcher)>([&](const WhereNode* where,Matcher* ctx) {`
`2059`	`2059`	`auto Temp = getResultAccesses(where->producer).first[0];`
`2060`	`2060`	`if (Temp.getTensorVar() == ws) {`
`2061`	`2061`	`for (auto i :getIndexVars()){`
`2062`	`2062`	`TempVars.insert(i);`
`2063`	`2063`	`}`
`2064`	`2064`	`}`
	`2065`	`+ ctx->match(where->producer);`
	`2066`	`+ ctx->match(where->consumer);`
`2065`	`2067`	`}));`
`2066`	`2068`	`for (auto i : accels) {`
`2067`	`2069`	`if (TempVars.find(i) == TempVars.end()) {`