Skip to content

Commit

Permalink
Nested rfactor fix, update predicates as well as source. (apache#3382)
Browse files Browse the repository at this point in the history
* Nested rfactor fix, update predicates as well as source.

* Linter

* Syntax fix.
  • Loading branch information
csarofeen authored and Wei Chen committed Jun 28, 2019
1 parent 4bd867b commit 5379923
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions src/schedule/schedule_dataflow_rewrite.cc
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
*
* http://www.apache.org/licenses/LICENSE-2.0
*
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
Expand Down Expand Up @@ -761,12 +761,15 @@ Array<Tensor> Schedule::rfactor(const Tensor& tensor,
VarReplacer replacer(vsub);
Array<Expr> new_source = ir::UpdateArray(reduce->source,
[&replacer] (const Expr& e) { return replacer.Mutate(e); });

Expr new_pred = replacer.Mutate(predicate);

std::vector<Expr> body;
for (size_t idx = 0; idx < reduce->source.size(); ++idx) {
body.emplace_back(Reduce::make(reduce->combiner,
new_source,
n->reduce_axis,
predicate,
new_pred,
idx));
}
n->body = Array<Expr>(body);
Expand Down

0 comments on commit 5379923

Please sign in to comment.