make_reduction.h

Go to the documentation of this file.

#ifndef FREE_TENSOR_MAKE_REDUCTION_H
#define FREE_TENSOR_MAKE_REDUCTION_H
 
#include <optional>
#include <unordered_set>
 
#include <func.h>
#include <mutator.h>
 
namespace freetensor {
 
class MakeReduction : public Mutator {
    const std::unordered_set<ReduceOp> &types_;
    bool canonicalOnly_;
 
  public:
    MakeReduction(const std::unordered_set<ReduceOp> &types, bool canonicalOnly)
        : types_(types), canonicalOnly_(canonicalOnly) {}
 
  private:
    bool isSameElem(const Store &s, const Load &l);
 
    Stmt doMake(Store op, ASTNodeType binOp, ReduceOp reduceOp,
                std::optional<ASTNodeType> invBinOp = std::nullopt);
 
  protected:
    Stmt visit(const Store &op) override;
};
 
inline Stmt makeReduction(const Stmt &op,
                          const std::unordered_set<ReduceOp> &types,
                          bool canonicalOnly = false) {
    return MakeReduction(types, canonicalOnly)(op);
}
 
inline Stmt makeReduction(const Stmt &op) {
    return makeReduction(op, {ReduceOp::Add, ReduceOp::Mul, ReduceOp::Min,
                              ReduceOp::Max, ReduceOp::LAnd, ReduceOp::LOr});
}
 
DEFINE_PASS_FOR_FUNC(makeReduction)
 
} // namespace freetensor
 
#endif // FREE_TENSOR_MAKE_REDUCTION_H