Commit 8fe715fe by Christian Sarofeen Committed by Lianmin Zheng

Nested rfactor fix, update predicates as well as source. (#3382)

* Nested rfactor fix, update predicates as well as source.

* Linter

* Syntax fix.
parent 8ef22176
...@@ -6,9 +6,9 @@ ...@@ -6,9 +6,9 @@
* to you under the Apache License, Version 2.0 (the * to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance * "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at * with the License. You may obtain a copy of the License at
* *
* http://www.apache.org/licenses/LICENSE-2.0 * http://www.apache.org/licenses/LICENSE-2.0
* *
* Unless required by applicable law or agreed to in writing, * Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an * software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
...@@ -761,12 +761,15 @@ Array<Tensor> Schedule::rfactor(const Tensor& tensor, ...@@ -761,12 +761,15 @@ Array<Tensor> Schedule::rfactor(const Tensor& tensor,
VarReplacer replacer(vsub); VarReplacer replacer(vsub);
Array<Expr> new_source = ir::UpdateArray(reduce->source, Array<Expr> new_source = ir::UpdateArray(reduce->source,
[&replacer] (const Expr& e) { return replacer.Mutate(e); }); [&replacer] (const Expr& e) { return replacer.Mutate(e); });
Expr new_pred = replacer.Mutate(predicate);
std::vector<Expr> body; std::vector<Expr> body;
for (size_t idx = 0; idx < reduce->source.size(); ++idx) { for (size_t idx = 0; idx < reduce->source.size(); ++idx) {
body.emplace_back(Reduce::make(reduce->combiner, body.emplace_back(Reduce::make(reduce->combiner,
new_source, new_source,
n->reduce_axis, n->reduce_axis,
predicate, new_pred,
idx)); idx));
} }
n->body = Array<Expr>(body); n->body = Array<Expr>(body);
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment