tlopex commented on code in PR #19667:
URL: https://github.com/apache/tvm/pull/19667#discussion_r3352690168
##########
src/target/z3/z3_prover_on.cc:
##########
@@ -0,0 +1,788 @@
+#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/extra/structural_equal.h>
+#include <tvm/ffi/extra/structural_hash.h>
+#include <tvm/runtime/logging.h>
+#include <tvm/tirx/analysis.h>
+#include <tvm/tirx/builtin.h>
+#include <tvm/tirx/expr.h>
+#include <tvm/tirx/expr_functor.h>
+#include <tvm/tirx/op.h>
+#include <tvm/tirx/op_attr_types.h>
+
+#include <algorithm>
+#include <climits>
+#include <map>
+#include <memory>
+#include <optional>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+#include "tvm/arith/analyzer.h"
+#include "tvm/ffi/cast.h"
+#include "tvm/ffi/object.h"
+#include "tvm/ffi/string.h"
+#include "tvm/ir/expr.h"
+#include "tvm/runtime/data_type.h"
+#include "z3++.h"
+
+namespace tvm::arith {
+
+using namespace tirx;
+using namespace ffi;
+
+namespace {
+
+struct Namespace {
+ std::unordered_set<std::string> used_names;
+ /// @brief Get a new name that is not used before
+ /// This function is used to generate z3 variable names
+ ///
+ /// Z3 may deduplicate variables with the same name, which
+ /// causes issues when different TVM variables are mapped to
+ /// the same z3 variable.
+ ///
+ /// This function generates unique names by appending
+ /// suffixes to the original expression string representation.
+ ///
+ /// such as : "x", "x$1", "x$2", ...
+ std::string GetNewName(const PrimExpr& expr) {
+ std::stringstream ss;
+ ss << expr;
+ auto name = ss.str();
+ if (used_names.count(name) == 0) {
+ used_names.insert(name);
+ return name;
+ }
+ int idx = 1;
+ std::string check_name = name + "$" + std::to_string(idx);
+ while (used_names.count(check_name)) {
+ idx++;
+ check_name = name + "$" + std::to_string(idx);
+ }
+ used_names.insert(check_name);
+ return check_name;
+ }
+};
+
+} // namespace
+
+class Z3Prover::Impl : ExprFunctor<z3::expr(const PrimExpr&)> {
+ public:
+ using Base = ExprFunctor<z3::expr(const PrimExpr&)>;
+ using Self = Z3Prover::Impl;
+
+ Analyzer* analyzer;
+ /// @brief Z3 context, a shared ptr, because tilelang want to copy the
Analyzer
+ // We use a thread_local static Z3 context so all analyzers within the same
thread
+ // can share a common context, because Z3 initialization is slow on some CPUs
+ // (e.g., AMD EPYC 7502 32-Core). Using thread_local ensures thread safety.
+ inline static thread_local std::shared_ptr<z3::context> ctx{new
z3::context()};
+
+ /// @brief Z3 solver instance
+ z3::solver solver{*ctx};
+
+ /// @brief Memorize pure expressions
+ std::unordered_map<PrimExpr, z3::expr, StructuralHash, ExprDeepEqual> memo_;
+
+ bool is_assume = false;
+
+ /// @brief Namespace for variable naming
+ Namespace ns;
+
+ /// @brief Timeout in milliseconds
+ unsigned timeout_ms{UINT_MAX};
+
+ /// @brief Max steps
+ unsigned rlimit{UINT_MAX};
+
+ /// @brief Create a z3 solver with custom options
+ static z3::solver CreateSolver(z3::context& ctx) {
+ z3::solver solver(ctx);
+ // here we disable model generation to speed up the solving process
+ solver.set("model", false);
+ // ensure determinstic behavior
+ solver.set("random_seed", (unsigned)42);
+ return solver;
+ }
+
+ Impl(Analyzer* parent) : analyzer(parent) {
+ scope_stack_.push_back({});
+ solver = CreateSolver(*ctx);
+ // default timeout 5ms
+ // Z3's implementation of timeout, when setting timeout T ms, it will stop
at T - 1 ms
+ // SetTimeoutMs(5);
+ // use rlimit, not timeout to ensure determinstic behavior
+ SetRLimit(1e4);
+ }
+
+ /// @brief Create a Free z3 expression from PrimExprNode
+ z3::expr Create(const PrimExprNode* op) {
+ auto ref = ffi::GetRef<PrimExpr>(op);
+ auto dtype = op->dtype;
+ std::string name = ns.GetNewName(ref);
+ /// TVM max_val can't handle uint64 max correctly, so we special case it
here
+ if (dtype.is_bool()) {
+ return ctx->bool_const(name.c_str());
+ } else {
+ z3::expr e = ctx->int_const(name.c_str());
+ if (dtype.is_uint() && dtype.bits() == 64) {
+ solver.add(ctx->int_val(0) <= e && e <=
ctx->int_val((uint64_t)UINT64_MAX));
+ } else {
+ auto min_val = Downcast<IntImm>(min_value(dtype))->value;
+ auto max_val = Downcast<IntImm>(max_value(dtype))->value;
+ solver.add(ctx->int_val(min_val) <= e && e <= ctx->int_val(max_val));
+ }
+ return e;
+ }
+ }
+
+ struct Scope {
+ enum Kind {
+ BindValue,
+ BindRange,
+ Constraint,
+ } kind;
+ Var var;
+ PrimExpr value;
+ PrimExpr min;
+ PrimExpr extent;
+ PrimExpr constraint;
+ };
+
+ /// @brief scope_stack memorizes existing constraint and bindings
+ /// to generate SMTLIB2 representation with comments
+ std::vector<std::vector<Scope>> scope_stack_;
+
+ /// @brief Enter a constraint scope
+ std::function<void()> EnterConstraint(const PrimExpr& constraint, bool
is_assume = false) {
+ scope_stack_.push_back({});
+ scope_stack_.back().push_back(
+ Scope{Scope::Constraint, Var(), PrimExpr(), PrimExpr(), PrimExpr(),
constraint});
+ solver.push();
+ this->is_assume = is_assume;
+ solver.add(VisitBool(constraint));
+ this->is_assume = false;
+ auto side_effect_exprs = std::move(side_effect_exprs_);
+ side_effect_exprs_.clear();
+ if (is_assume) {
+ return [this, side_effect_exprs]() {
+ solver.pop();
+ for (const auto& expr : side_effect_exprs) {
+ memo_.erase(expr);
+ }
+ scope_stack_.pop_back();
+ };
+ } else {
+ for (const auto& expr : side_effect_exprs) {
+ memo_.erase(expr);
+ }
+ return [this]() {
+ solver.pop();
+ scope_stack_.pop_back();
+ };
+ }
+ }
+
+ /// @brief Check trivil bad cases, return true if the expr is a bad case
+ /// Z3 prover may take a long time to initialize (at least 200us),
+ /// This optimization can speedup 30% of the test cases in our unit tests
+ bool CheckTrivilBadCases(const PrimExpr& expr) {
+ if (IsFreeNode(expr)) {
+ return true;
+ }
+ auto checkTrivilCmp = [this](const PrimExpr& lhs, const PrimExpr& rhs) {
+ if (IsFreeNode(lhs) && rhs->IsInstance<IntImmNode>()) {
+ return true;
+ }
+ if (IsFreeNode(rhs) && lhs->IsInstance<IntImmNode>()) {
+ return true;
+ }
+ if (IsFreeNode(lhs) && IsFreeNode(rhs)) {
+ return true;
+ }
+ // cast('xxx', free_var) == constant
+ if (auto cast = lhs.as<CastNode>()) {
+ if (IsFreeNode(cast->value) && rhs->IsInstance<IntImmNode>()) {
+ return true;
+ }
+ }
+ // constant == cast('xxx', free_var)
+ if (auto cast = rhs.as<CastNode>()) {
+ if (IsFreeNode(cast->value) && lhs->IsInstance<IntImmNode>()) {
+ return true;
+ }
+ }
+ return false;
+ };
+ if (auto eq = expr.as<EQNode>()) {
+ auto lhs = eq->a;
+ auto rhs = eq->b;
+ return checkTrivilCmp(lhs, rhs);
+ } else if (auto ne = expr.as<NENode>()) {
+ auto lhs = ne->a;
+ auto rhs = ne->b;
+ return checkTrivilCmp(lhs, rhs);
+ }
+ return false;
+ }
+
+ /// @brief Check if the expression can be proved
+ bool CanProve(const PrimExpr& expr) {
+ if (CheckTrivilBadCases(expr)) return false;
+ if (!IsValidDType(expr->dtype)) return false;
+ z3::expr_vector constr(*ctx);
+ constr.push_back(!ConvertBool(expr));
+ auto result = solver.check(constr);
+ constr.pop_back();
+ return result == z3::unsat;
+ }
+
+ /// @brief Binded
+ /// @brief Bind a variable to a value or a range
+ void Bind(const Var& var, const PrimExpr& value, bool allow_override =
false) {
+ if (!IsValidDType(var->dtype)) return;
+ scope_stack_.back().push_back(Scope{Scope::BindValue, var, value});
+ // we add the binding whenever the value is pure,
+ // because non-pure parts are handling by creating free variables in
VisitExpr
+ memo_.emplace(var, ConvertInt(value));
+ }
+
+ /// @brief Bind a variable to a range
+ void Bind(const Var& var, const Range& range, bool allow_override = false) {
+ if (!IsValidDType(var->dtype)) return;
+ scope_stack_.back().push_back(
+ Scope{Scope::BindRange, var, PrimExpr(), range->min, range->extent});
+ // 1. Create a placeholder for the var, and save it in the memo
+ // if the var is overrided later, we can just update the memo, and the
old placeholder will
+ // be ignored
+ auto var_expr = Create(var.as<PrimExprNode>());
+ memo_.emplace(var, var_expr);
+
+ // 2. Add constraint on the placeholder
+ // when min_expr >= max_expr, the range is empty, which is under
undefined behavior
+ // instead of adding an unsat constraint, we just skip the range
constraint to leave it a
+ // free var
+ if (tirx::is_const_int(range->min) && tirx::is_const_int(range->min +
range->extent)) {
+ int64_t min_value = *tirx::as_const_int(range->min);
+ int64_t max_value = *tirx::as_const_int(range->min + range->extent);
+ if (min_value < max_value) {
+ solver.add(ctx->int_val(min_value) <= var_expr);
+ solver.add(var_expr < ctx->int_val(max_value));
+ }
+ } else {
+ solver.add(ConvertBool(range->extent <= 0 ||
+ (range->min <= var && var < range->min +
range->extent)));
+ }
+ }
+
+ void CopyFrom(const Self& other_) {
+ // 1. create a new solver
+ // because this->solver depends on this->ctx
+ // we need to deconstruct the old solver, and create a new one
depending on other_.ctx
+ solver = CreateSolver(*other_.ctx);
+ // 2. copy the context
+ // the context is a shared_ptr, we can just copy the pointer
+ ctx = other_.ctx;
Review Comment:
I think we can remove it as gemini suggested
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
