int_solver.h

Go to the documentation of this file.

/*
 * 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.
 */

#ifndef TVM_ARITH_INT_SOLVER_H_
#define TVM_ARITH_INT_SOLVER_H_

#include <tvm/ir/expr.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/op.h>

#include <unordered_map>
#include <utility>
#include <vector>

#include "analyzer.h"

namespace tvm {
namespace arith {

using tir::IterVar;
using tir::Var;
using tir::VarNode;

// According to experiments two best simplifications orders were can->rw and rw->can->rw,
// but rw->can->rw is better for a couple of cases.
// Also we should end with rw because it factors multipliers out.
constexpr int kSimplifyRewriteCanonicalRewrite = 3;

class IntGroupBoundsNode : public Object {
 public:
  PrimExpr coef;
  Array<PrimExpr> lower;
  Array<PrimExpr> equal;
  Array<PrimExpr> upper;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("coef", &coef);
    v->Visit("lower", &lower);
    v->Visit("equal", &equal);
    v->Visit("upper", &upper);
  }

  bool SEqualReduce(const IntGroupBoundsNode* other, SEqualReducer eq) const {
    return eq(coef, other->coef) && eq(lower, other->lower) && eq(equal, other->equal) &&
           eq(upper, other->upper);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(coef);
    hash_reduce(lower);
    hash_reduce(equal);
    hash_reduce(upper);
  }

  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const char* _type_key = "arith.IntGroupBounds";
  TVM_DECLARE_FINAL_OBJECT_INFO(IntGroupBoundsNode, Object);
};

class IntGroupBounds : public ObjectRef {
 public:
  TVM_DLL IntGroupBounds(PrimExpr coef, Array<PrimExpr> lower, Array<PrimExpr> equal,
                         Array<PrimExpr> upper);

  static IntGroupBounds FromRange(const Range& r);

  IntGroupBounds Substitute(const Map<Var, PrimExpr>& subst) const;

  Range FindBestRange(const Map<Var, Range>& vranges_addl = {}) const;

  IntGroupBounds operator+(const Range& r);

  TVM_DEFINE_OBJECT_REF_METHODS(IntGroupBounds, ObjectRef, IntGroupBoundsNode);
};

class IntConstraintsNode : public Object {
 public:
  // e.g., \alpha, \beta, must be integers
  Array<Var> variables;
  // e.g., 1 <= \alpha <= N, etc.
  // it is absolutely ok to include ranges for parameters
  // (variables that are not in this->variables) in this map
  Map<Var, Range> ranges;
  // linear equalities or inequalities
  // e.g., A \alpha = \beta or A \alpha <= \beta
  Array<PrimExpr> relations;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("variables", &variables);
    v->Visit("ranges", &ranges);
    v->Visit("relations", &relations);
  }

  bool SEqualReduce(const IntConstraintsNode* other, SEqualReducer equal) const {
    return equal(variables, other->variables) && equal(ranges, other->ranges) &&
           equal(relations, other->relations);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(variables);
    hash_reduce(ranges);
    hash_reduce(relations);
  }

  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const char* _type_key = "arith.IntConstraints";
  TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsNode, Object);
};

class IntConstraints : public ObjectRef {
 public:
  TVM_DLL IntConstraints(Array<Var> variables, Map<Var, Range> ranges, Array<PrimExpr> relations);

  TVM_DEFINE_OBJECT_REF_METHODS(IntConstraints, ObjectRef, IntConstraintsNode);
};

class IntConstraintsTransformNode : public Object {
 public:
  IntConstraints src;
  IntConstraints dst;
  Map<Var, PrimExpr> src_to_dst;
  Map<Var, PrimExpr> dst_to_src;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("src", &src);
    v->Visit("dst", &dst);
    v->Visit("src_to_dst", &src_to_dst);
    v->Visit("dst_to_src", &dst_to_src);
  }

  bool SEqualReduce(const IntConstraintsTransformNode* other, SEqualReducer equal) const {
    return equal(src, other->src) && equal(dst, other->dst) &&
           equal(src_to_dst, other->src_to_dst) && equal(dst_to_src, other->dst_to_src);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(src);
    hash_reduce(dst);
    hash_reduce(src_to_dst);
    hash_reduce(dst_to_src);
  }

  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const char* _type_key = "arith.IntConstraintsTransform";
  TVM_DECLARE_FINAL_OBJECT_INFO(IntConstraintsTransformNode, Object);
};

class IntConstraintsTransform : public ObjectRef {
 public:
  TVM_DLL IntConstraintsTransform(IntConstraints src, IntConstraints dst,
                                  Map<Var, PrimExpr> src_to_dst, Map<Var, PrimExpr> dst_to_src);

  IntConstraintsTransform operator+(const IntConstraintsTransform& other) const;

  TVM_DEFINE_OBJECT_REF_METHODS(IntConstraintsTransform, ObjectRef, IntConstraintsTransformNode);
};

typedef std::pair<Map<Var, IntGroupBounds>, Array<PrimExpr>> PartialSolvedInequalities;

void SmithNormalFormDiag(std::vector<std::vector<int64_t>>* S, std::vector<std::vector<int64_t>>* V,
                         std::vector<PrimExpr>* x, std::vector<PrimExpr>* y);

IntConstraintsTransform SolveLinearEquations(const IntConstraints& system_to_solve);

PartialSolvedInequalities SolveLinearInequalities(const IntConstraints& system_to_solve);

Array<PrimExpr> AsConditions(const Array<Var>& variables, const Map<Var, IntGroupBounds>& bounds,
                             const Array<PrimExpr>& relations);

IntConstraints SolveInequalitiesToRange(const IntConstraints& system_to_solve);

IntConstraintsTransform SolveInequalitiesDeskewRange(const IntConstraints& system_to_solve);

}  // namespace arith
}  // namespace tvm
#endif  // TVM_ARITH_INT_SOLVER_H_