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/ir/prim/expr.h>
#include <tvm/tirx/op.h>
 
#include <unordered_map>
#include <utility>
#include <vector>
 
#include "analyzer.h"
 
namespace tvm {
namespace arith {
 
using tirx::IterVar;
using tirx::PrimVar;
using tirx::Var;
using tirx::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 ffi::Object {
 public:
  PrimExpr coef;
  ffi::Array<PrimExpr> lower;
  ffi::Array<PrimExpr> equal;
  ffi::Array<PrimExpr> upper;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IntGroupBoundsNode>()
        .def_ro("coef", &IntGroupBoundsNode::coef)
        .def_ro("lower", &IntGroupBoundsNode::lower)
        .def_ro("equal", &IntGroupBoundsNode::equal)
        .def_ro("upper", &IntGroupBoundsNode::upper);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntGroupBounds", IntGroupBoundsNode, ffi::Object);
};
 
class IntGroupBounds : public ffi::ObjectRef {
 public:
  TVM_DLL IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower, ffi::Array<PrimExpr> equal,
                         ffi::Array<PrimExpr> upper);
 
  static IntGroupBounds FromRange(const Range& r);
 
  IntGroupBounds Substitute(const ffi::Map<Var, PrimExpr>& subst) const;
 
  Range FindBestRange(const ffi::Map<Var, Range>& vranges_addl = {}) const;
 
  IntGroupBounds operator+(const Range& r);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntGroupBounds, ffi::ObjectRef, IntGroupBoundsNode);
};
 
class IntConstraintsNode : public ffi::Object {
 public:
  // e.g., \alpha, \beta, must be integers
  ffi::Array<PrimVar> 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
  ffi::Map<Var, Range> ranges;
  // linear equalities or inequalities
  // e.g., A \alpha = \beta or A \alpha <= \beta
  ffi::Array<PrimExpr> relations;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IntConstraintsNode>()
        .def_ro("variables", &IntConstraintsNode::variables)
        .def_ro("ranges", &IntConstraintsNode::ranges)
        .def_ro("relations", &IntConstraintsNode::relations);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraints", IntConstraintsNode, ffi::Object);
};
 
class IntConstraints : public ffi::ObjectRef {
 public:
  TVM_DLL IntConstraints(ffi::Array<PrimVar> variables, ffi::Map<Var, Range> ranges,
                         ffi::Array<PrimExpr> relations);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraints, ffi::ObjectRef, IntConstraintsNode);
};
 
class IntConstraintsTransformNode : public ffi::Object {
 public:
  IntConstraints src;
  IntConstraints dst;
  ffi::Map<Var, PrimExpr> src_to_dst;
  ffi::Map<Var, PrimExpr> dst_to_src;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IntConstraintsTransformNode>()
        .def_ro("src", &IntConstraintsTransformNode::src)
        .def_ro("dst", &IntConstraintsTransformNode::dst)
        .def_ro("src_to_dst", &IntConstraintsTransformNode::src_to_dst)
        .def_ro("dst_to_src", &IntConstraintsTransformNode::dst_to_src);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IntConstraintsTransform", IntConstraintsTransformNode,
                                    ffi::Object);
};
 
class IntConstraintsTransform : public ffi::ObjectRef {
 public:
  TVM_DLL IntConstraintsTransform(IntConstraints src, IntConstraints dst,
                                  ffi::Map<Var, PrimExpr> src_to_dst,
                                  ffi::Map<Var, PrimExpr> dst_to_src);
 
  IntConstraintsTransform operator+(const IntConstraintsTransform& other) const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntConstraintsTransform, ffi::ObjectRef,
                                             IntConstraintsTransformNode);
};
 
typedef std::pair<ffi::Map<Var, IntGroupBounds>, ffi::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);
 
ffi::Array<PrimExpr> AsConditions(const ffi::Array<PrimVar>& variables,
                                  const ffi::Map<Var, IntGroupBounds>& bounds,
                                  const ffi::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_