src/arith/modular_set.cc - tvm - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * 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
  * KIND, either express or implied.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 /*!
  * \file modular_set.cc
  * \brief Modular set analysis
  */
 #include <tvm/arith/analyzer.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/tir/builtin.h>
 #include <tvm/tir/expr_functor.h>
 #include <tvm/tir/op.h>

 #include <limits>
 #include <unordered_map>
 #include <utility>

 #include "pattern_match.h"

 namespace tvm {
 namespace arith {

 using namespace tir;

 TVM_REGISTER_NODE_TYPE(ModularSetNode);

 ModularSet::ModularSet(int64_t coeff, int64_t base) {
   auto node = make_object<ModularSetNode>();
   node->coeff = coeff;
   node->base = base;
   // finish construction.
   data_ = std::move(node);
 }

 TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
     .set_dispatch<ModularSetNode>([](const ObjectRef& node, ReprPrinter* p) {
       auto* op = static_cast<const ModularSetNode*>(node.get());
       p->stream << "ModularSet("
                 << "coeff=" << op->coeff << ", base=" << op->base << ')';
     });

 ModularSet MakeModularSet(int64_t coeff, int64_t base) { return ModularSet(coeff, base); }

 TVM_REGISTER_GLOBAL("arith.ModularSet").set_body_typed(MakeModularSet);

 // internal entry for const int bound
 struct ModularSetAnalyzer::Entry {
   int64_t coeff{1};
   int64_t base{0};

   Entry() = default;

   Entry(int64_t coeff, int64_t base) {
     CHECK_GE(coeff, 0);
     this->coeff = coeff;
     if (coeff != 0) {
       base = base % coeff;
       if (base < 0) base += coeff;
     }
     this->base = base;
   }

   bool is_const() const { return coeff == 0; }

   bool operator==(const Entry& other) const { return coeff == other.coeff && base == other.base; }

   bool operator==(const ModularSet& other) const {
     return other.defined() && coeff == other->coeff && base == other->base;
   }
 };

 class ModularSetAnalyzer::Impl : public ExprFunctor<ModularSetAnalyzer::Entry(const PrimExpr&)> {
  public:
   explicit Impl(Analyzer* parent) : parent_(parent) {}

   void Update(const Var& var, const ModularSet& info, bool allow_override) {
     if (!allow_override) {
       auto it = var_map_.find(var);
       if (it != var_map_.end()) {
         CHECK(it->second == info) << "Trying to update var \'" << var << "\'"
                                   << " with a different const bound: "
                                   << "original=" << ModularSet(it->second.coeff, it->second.base)
                                   << ", new=" << info;
       }
     }
     var_map_[var] = Entry(info->coeff, info->base);
   }

   // Detect useful constraints and use them in the analysis scope.
   std::function<void()> EnterConstraint(const PrimExpr& constraint) {
     PVar<Var> var;
     PVar<IntImm> coeff, base;
     // pattern match interesting constraints
     if ((truncmod(var, coeff) == base).Match(constraint) ||
         (floormod(var, coeff) == base).Match(constraint)) {
       Entry entry(coeff.Eval()->value, base.Eval()->value);
       return UpdateByIntersect(var.Eval(), entry);
     }
     return nullptr;
   }

   // Override visitor behaviors
   Entry VisitExprDefault_(const Object* op) final { return Everything(); }

   Entry VisitExpr_(const LetNode* op) final {
     auto it = var_map_.find(op->var);
     // if the var has not been binded, update the info.
     if (it == var_map_.end()) {
       var_map_[op->var] = this->VisitExpr(op->value);
       Entry ret = VisitExpr(op->body);
       var_map_.erase(op->var);
       return ret;
     } else {
       return VisitExpr(op->body);
     }
   }

   Entry VisitExpr_(const CastNode* op) final { return VisitExpr(op->value); }

   Entry VisitExpr_(const IntImmNode* op) final { return Entry(0, op->value); }

   Entry VisitExpr_(const AddNode* op) final {
     Entry a = VisitExpr(op->a);
     Entry b = VisitExpr(op->b);
     int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
     return Entry(coeff, a.base + b.base);
   }

   Entry VisitExpr_(const SubNode* op) final {
     Entry a = VisitExpr(op->a);
     Entry b = VisitExpr(op->b);
     int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
     return Entry(coeff, a.base - b.base);
   }

   Entry VisitExpr_(const MulNode* op) final {
     Entry a = VisitExpr(op->a);
     Entry b = VisitExpr(op->b);
     // Simplification rule, x, y, z are in Z
     // (p x + n) (q y + m)
     // -> pq xy + pm x + qn y + mn
     // -> pq z + pm x + qn y + mn
     int64_t pq = a.coeff * b.coeff;
     int64_t pm = a.coeff * b.base;
     int64_t qn = a.base * b.coeff;
     int64_t coeff = ZeroAwareGCD(pq, ZeroAwareGCD(pm, qn));
     return Entry(coeff, a.base * b.base);
   }

   Entry DivByConst(const PrimExpr& lhs, int64_t val, bool round_down) {
     Entry a = VisitExpr(lhs);
     CHECK_NE(val, 0);
     if (a.coeff % val == 0) {
       if (a.base == 0) {
         // a c x  / c -> a x
         return Entry(std::abs(a.coeff / val), 0);
       }
       // positive division have a clear rounding mode.
       // Only handle case where we clearly know we need to round down.
       if (a.base > 0 && val > 0 && (round_down || parent_->CanProveGreaterEqual(lhs, 0))) {
         return Entry(a.coeff / val, a.base / val);
       }
     }
     return Everything();
   }

   Entry VisitExpr_(const DivNode* op) final {
     Entry b = VisitExpr(op->b);
     if (b.is_const()) {
       return DivByConst(op->a, b.base, false);
     }
     return Everything();
   }

   Entry VisitExpr_(const FloorDivNode* op) final {
     Entry b = VisitExpr(op->b);
     if (b.is_const()) {
       return DivByConst(op->a, b.base, true);
     }
     return Everything();
   }

   Entry VisitExpr_(const MinNode* op) final {
     Entry a = VisitExpr(op->a);
     Entry b = VisitExpr(op->b);
     return Union(a, b);
   }

   Entry VisitExpr_(const MaxNode* op) final {
     Entry a = VisitExpr(op->a);
     Entry b = VisitExpr(op->b);
     return Union(a, b);
   }

   Entry VisitExpr_(const SelectNode* op) final {
     Entry a = VisitExpr(op->true_value);
     Entry b = VisitExpr(op->false_value);
     return Union(a, b);
   }

   Entry VisitExpr_(const CallNode* op) final {
     // only special handle >> which can be
     // used for index calculation.
     if (op->op.same_as(tir::builtin::shift_right())) {
       return VisitRightShift(op);
     } else {
       return Everything();
     }
   }

   Entry VisitExpr_(const VarNode* op) final {
     Var v = GetRef<Var>(op);
     auto it = var_map_.find(v);
     if (it != var_map_.end()) {
       return it->second;
     } else {
       return Everything();
     }
   }

   Entry VisitRightShift(const CallNode* op) {
     Entry b = VisitExpr(op->args[1]);
     // a c x  / c -> a x
     if (b.is_const()) {
       return DivByConst(op->args[0], static_cast<int64_t>(1) << b.base, true);
     }
     return Everything();
   }

  private:
   /*! \brief pointer to parent. */
   Analyzer* parent_{nullptr};
   // internal variable map
   std::unordered_map<Var, Entry, ObjectPtrHash, ObjectPtrEqual> var_map_;
   /*!
    * \brief Update var by intersecting entry with var's current set.
    * \param var The variable.
    * \param entry The entry to be updated.
    * \return The recovery function of the scope.
    */
   std::function<void()> UpdateByIntersect(const Var& var, Entry entry) {
     Entry old = Everything();
     auto it = var_map_.find(var);
     if (it != var_map_.end()) {
       old = it->second;
     }
     var_map_[var] = Intersect(old, entry);
     // reover function.
     return [this, old, var]() { var_map_[var] = old; };
   }
   /*!
    * \brief Create union of two sets.
    * \param a The left operand.
    * \param b the right operand.
    */
   static Entry Union(Entry a, Entry b) {
     // {ax + y} \cup {bz + h} => {gcd(a, b) x + {y or h}}
     int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
     if (coeff == 0) {
       if (a.base == b.base) return a;
       return Everything();
     }
     int64_t base0 = a.base % coeff;
     int64_t base1 = b.base % coeff;
     if (base0 == base1) {
       return Entry(coeff, base0);
     } else {
       return Entry(ZeroAwareGCD(ZeroAwareGCD(base0, base1), coeff), base0);
     }
   }

   /*!
    * \brief Create interect of two sets.
    * \param a The left operand.
    * \param b the right operand.
    */
   static Entry Intersect(Entry a, Entry b) {
     int64_t x, y;
     int64_t c1 = a.coeff, b1 = a.base, c2 = b.coeff, b2 = b.base;
     // z = c1 * p + b1
     // z = c2 * q + b2
     // c1 * x + c2 * y = gcd(c1, c2)
     // -> c1 * p - c2 * q = b2 - b1
     // -> p = (b2 - b1) / gcd * x
     // -> q = (b2 - b1) / gcd * (-y)
     // -> z = LCM(x, y) * k + (c1 * p + b1)
     int64_t gcd = ExtendedEuclidean(c1, c2, &x, &y);
     int64_t v = b2 - b1;
     if (v % gcd == 0) {
       x = v / gcd * x;
       y = v / gcd * (-y);
       int64_t coeff = c1 / gcd * c2;
       return Entry(coeff, x * c1 + b1);
     } else {
       return Nothing();
     }
   }
   /*!
    * \brief return everything dtype can represent.
    * \return Bound that represent everything dtype can represent.
    */
   static Entry Everything() { return Entry(1, 0); }
   /*!
    * \brief return an empty set
    * \return Bound that represent everything dtype can represent.
    */
   static Entry Nothing() { return Entry(0, 1); }
 };

 ModularSet ModularSetAnalyzer::operator()(const PrimExpr& expr) {
   Entry ret = impl_->VisitExpr(expr);
   return ModularSet(ret.coeff, ret.base);
 }

 void ModularSetAnalyzer::Update(const Var& var, const ModularSet& info, bool allow_override) {
   impl_->Update(var, info, allow_override);
 }

 std::function<void()> ModularSetAnalyzer::EnterConstraint(const PrimExpr& constraint) {
   return impl_->EnterConstraint(constraint);
 }

 ModularSetAnalyzer::ModularSetAnalyzer(Analyzer* parent) : impl_(new Impl(parent)) {}

 ModularSetAnalyzer::~ModularSetAnalyzer() { delete impl_; }

 }  // namespace arith
 }  // namespace tvm
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* 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
	* KIND, either express or implied. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	/*!
	* \file modular_set.cc
	* \brief Modular set analysis
	*/
	#include <tvm/arith/analyzer.h>
	#include <tvm/runtime/registry.h>
	#include <tvm/tir/builtin.h>
	#include <tvm/tir/expr_functor.h>
	#include <tvm/tir/op.h>

	#include <limits>
	#include <unordered_map>
	#include <utility>

	#include "pattern_match.h"

	namespace tvm {
	namespace arith {

	using namespace tir;

	TVM_REGISTER_NODE_TYPE(ModularSetNode);

	ModularSet::ModularSet(int64_t coeff, int64_t base) {
	auto node = make_object<ModularSetNode>();
	node->coeff = coeff;
	node->base = base;
	// finish construction.
	data_ = std::move(node);
	}

	TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
	.set_dispatch<ModularSetNode>([](const ObjectRef& node, ReprPrinter* p) {
	auto* op = static_cast<const ModularSetNode*>(node.get());
	p->stream << "ModularSet("
	<< "coeff=" << op->coeff << ", base=" << op->base << ')';
	});

	ModularSet MakeModularSet(int64_t coeff, int64_t base) { return ModularSet(coeff, base); }

	TVM_REGISTER_GLOBAL("arith.ModularSet").set_body_typed(MakeModularSet);

	// internal entry for const int bound
	struct ModularSetAnalyzer::Entry {
	int64_t coeff{1};
	int64_t base{0};

	Entry() = default;

	Entry(int64_t coeff, int64_t base) {
	CHECK_GE(coeff, 0);
	this->coeff = coeff;
	if (coeff != 0) {
	base = base % coeff;
	if (base < 0) base += coeff;
	}
	this->base = base;
	}

	bool is_const() const { return coeff == 0; }

	bool operator==(const Entry& other) const { return coeff == other.coeff && base == other.base; }

	bool operator==(const ModularSet& other) const {
	return other.defined() && coeff == other->coeff && base == other->base;
	}
	};

	class ModularSetAnalyzer::Impl : public ExprFunctor<ModularSetAnalyzer::Entry(const PrimExpr&)> {
	public:
	explicit Impl(Analyzer* parent) : parent_(parent) {}

	void Update(const Var& var, const ModularSet& info, bool allow_override) {
	if (!allow_override) {
	auto it = var_map_.find(var);
	if (it != var_map_.end()) {
	CHECK(it->second == info) << "Trying to update var \'" << var << "\'"
	<< " with a different const bound: "
	<< "original=" << ModularSet(it->second.coeff, it->second.base)
	<< ", new=" << info;
	}
	}
	var_map_[var] = Entry(info->coeff, info->base);
	}

	// Detect useful constraints and use them in the analysis scope.
	std::function<void()> EnterConstraint(const PrimExpr& constraint) {
	PVar<Var> var;
	PVar<IntImm> coeff, base;
	// pattern match interesting constraints
	if ((truncmod(var, coeff) == base).Match(constraint) \|\|
	(floormod(var, coeff) == base).Match(constraint)) {
	Entry entry(coeff.Eval()->value, base.Eval()->value);
	return UpdateByIntersect(var.Eval(), entry);
	}
	return nullptr;
	}

	// Override visitor behaviors
	Entry VisitExprDefault_(const Object* op) final { return Everything(); }

	Entry VisitExpr_(const LetNode* op) final {
	auto it = var_map_.find(op->var);
	// if the var has not been binded, update the info.
	if (it == var_map_.end()) {
	var_map_[op->var] = this->VisitExpr(op->value);
	Entry ret = VisitExpr(op->body);
	var_map_.erase(op->var);
	return ret;
	} else {
	return VisitExpr(op->body);
	}
	}

	Entry VisitExpr_(const CastNode* op) final { return VisitExpr(op->value); }

	Entry VisitExpr_(const IntImmNode* op) final { return Entry(0, op->value); }

	Entry VisitExpr_(const AddNode* op) final {
	Entry a = VisitExpr(op->a);
	Entry b = VisitExpr(op->b);
	int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
	return Entry(coeff, a.base + b.base);
	}

	Entry VisitExpr_(const SubNode* op) final {
	Entry a = VisitExpr(op->a);
	Entry b = VisitExpr(op->b);
	int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
	return Entry(coeff, a.base - b.base);
	}

	Entry VisitExpr_(const MulNode* op) final {
	Entry a = VisitExpr(op->a);
	Entry b = VisitExpr(op->b);
	// Simplification rule, x, y, z are in Z
	// (p x + n) (q y + m)
	// -> pq xy + pm x + qn y + mn
	// -> pq z + pm x + qn y + mn
	int64_t pq = a.coeff * b.coeff;
	int64_t pm = a.coeff * b.base;
	int64_t qn = a.base * b.coeff;
	int64_t coeff = ZeroAwareGCD(pq, ZeroAwareGCD(pm, qn));
	return Entry(coeff, a.base * b.base);
	}

	Entry DivByConst(const PrimExpr& lhs, int64_t val, bool round_down) {
	Entry a = VisitExpr(lhs);
	CHECK_NE(val, 0);
	if (a.coeff % val == 0) {
	if (a.base == 0) {
	// a c x / c -> a x
	return Entry(std::abs(a.coeff / val), 0);
	}
	// positive division have a clear rounding mode.
	// Only handle case where we clearly know we need to round down.
	if (a.base > 0 && val > 0 && (round_down \|\| parent_->CanProveGreaterEqual(lhs, 0))) {
	return Entry(a.coeff / val, a.base / val);
	}
	}
	return Everything();
	}

	Entry VisitExpr_(const DivNode* op) final {
	Entry b = VisitExpr(op->b);
	if (b.is_const()) {
	return DivByConst(op->a, b.base, false);
	}
	return Everything();
	}

	Entry VisitExpr_(const FloorDivNode* op) final {
	Entry b = VisitExpr(op->b);
	if (b.is_const()) {
	return DivByConst(op->a, b.base, true);
	}
	return Everything();
	}

	Entry VisitExpr_(const MinNode* op) final {
	Entry a = VisitExpr(op->a);
	Entry b = VisitExpr(op->b);
	return Union(a, b);
	}

	Entry VisitExpr_(const MaxNode* op) final {
	Entry a = VisitExpr(op->a);
	Entry b = VisitExpr(op->b);
	return Union(a, b);
	}

	Entry VisitExpr_(const SelectNode* op) final {
	Entry a = VisitExpr(op->true_value);
	Entry b = VisitExpr(op->false_value);
	return Union(a, b);
	}

	Entry VisitExpr_(const CallNode* op) final {
	// only special handle >> which can be
	// used for index calculation.
	if (op->op.same_as(tir::builtin::shift_right())) {
	return VisitRightShift(op);
	} else {
	return Everything();
	}
	}

	Entry VisitExpr_(const VarNode* op) final {
	Var v = GetRef<Var>(op);
	auto it = var_map_.find(v);
	if (it != var_map_.end()) {
	return it->second;
	} else {
	return Everything();
	}
	}

	Entry VisitRightShift(const CallNode* op) {
	Entry b = VisitExpr(op->args[1]);
	// a c x / c -> a x
	if (b.is_const()) {
	return DivByConst(op->args[0], static_cast<int64_t>(1) << b.base, true);
	}
	return Everything();
	}

	private:
	/! \brief pointer to parent. /
	Analyzer* parent_{nullptr};
	// internal variable map
	std::unordered_map<Var, Entry, ObjectPtrHash, ObjectPtrEqual> var_map_;
	/*!
	* \brief Update var by intersecting entry with var's current set.
	* \param var The variable.
	* \param entry The entry to be updated.
	* \return The recovery function of the scope.
	*/
	std::function<void()> UpdateByIntersect(const Var& var, Entry entry) {
	Entry old = Everything();
	auto it = var_map_.find(var);
	if (it != var_map_.end()) {
	old = it->second;
	}
	var_map_[var] = Intersect(old, entry);
	// reover function.
	return [this, old, var]() { var_map_[var] = old; };
	}
	/*!
	* \brief Create union of two sets.
	* \param a The left operand.
	* \param b the right operand.
	*/
	static Entry Union(Entry a, Entry b) {
	// {ax + y} \cup {bz + h} => {gcd(a, b) x + {y or h}}
	int64_t coeff = ZeroAwareGCD(a.coeff, b.coeff);
	if (coeff == 0) {
	if (a.base == b.base) return a;
	return Everything();
	}
	int64_t base0 = a.base % coeff;
	int64_t base1 = b.base % coeff;
	if (base0 == base1) {
	return Entry(coeff, base0);
	} else {
	return Entry(ZeroAwareGCD(ZeroAwareGCD(base0, base1), coeff), base0);
	}
	}

	/*!
	* \brief Create interect of two sets.
	* \param a The left operand.
	* \param b the right operand.
	*/
	static Entry Intersect(Entry a, Entry b) {
	int64_t x, y;
	int64_t c1 = a.coeff, b1 = a.base, c2 = b.coeff, b2 = b.base;
	// z = c1 * p + b1
	// z = c2 * q + b2
	// c1 * x + c2 * y = gcd(c1, c2)
	// -> c1 * p - c2 * q = b2 - b1
	// -> p = (b2 - b1) / gcd * x
	// -> q = (b2 - b1) / gcd * (-y)
	// -> z = LCM(x, y) * k + (c1 * p + b1)
	int64_t gcd = ExtendedEuclidean(c1, c2, &x, &y);
	int64_t v = b2 - b1;
	if (v % gcd == 0) {
	x = v / gcd * x;
	y = v / gcd * (-y);
	int64_t coeff = c1 / gcd * c2;
	return Entry(coeff, x * c1 + b1);
	} else {
	return Nothing();
	}
	}
	/*!
	* \brief return everything dtype can represent.
	* \return Bound that represent everything dtype can represent.
	*/
	static Entry Everything() { return Entry(1, 0); }
	/*!
	* \brief return an empty set
	* \return Bound that represent everything dtype can represent.
	*/
	static Entry Nothing() { return Entry(0, 1); }
	};

	ModularSet ModularSetAnalyzer::operator()(const PrimExpr& expr) {
	Entry ret = impl_->VisitExpr(expr);
	return ModularSet(ret.coeff, ret.base);
	}

	void ModularSetAnalyzer::Update(const Var& var, const ModularSet& info, bool allow_override) {
	impl_->Update(var, info, allow_override);
	}

	std::function<void()> ModularSetAnalyzer::EnterConstraint(const PrimExpr& constraint) {
	return impl_->EnterConstraint(constraint);
	}

	ModularSetAnalyzer::ModularSetAnalyzer(Analyzer* parent) : impl_(new Impl(parent)) {}

	ModularSetAnalyzer::~ModularSetAnalyzer() { delete impl_; }

	} // namespace arith
	} // namespace tvm