iter_affine_map.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_ITER_AFFINE_MAP_H_
#define TVM_ARITH_ITER_AFFINE_MAP_H_
 
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/cow.h>
#include <tvm/ir/expr.h>
#include <tvm/tirx/var.h>
 
namespace tvm {
namespace arith {
 
class IterMapExprNode : public ExprNode {
 public:
  static constexpr const uint32_t _type_child_slots = 2;
  TVM_FFI_DECLARE_OBJECT_INFO("arith.IterMapExpr", IterMapExprNode, ExprNode);
};
 
class IterMapExpr : public PrimExpr {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMapExpr, PrimExpr, IterMapExprNode);
  static constexpr bool _type_container_is_exact = true;
};
 
class IterMarkNode : public ffi::Object {
 public:
  PrimExpr source;
  PrimExpr extent;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IterMarkNode>()
        .def_ro("source", &IterMarkNode::source)
        .def_ro("extent", &IterMarkNode::extent);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindDAGNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMark", IterMarkNode, ffi::Object);
};
 
class IterMark : public ffi::ObjectRef {
 public:
  TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterMark, ffi::ObjectRef, IterMarkNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode);
};
 
class IterSplitExprNode : public IterMapExprNode {
 public:
  IterMark source;
  PrimExpr lower_factor;
  PrimExpr extent;
  PrimExpr scale;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IterSplitExprNode>()
        .def_ro("source", &IterSplitExprNode::source)
        .def_ro("lower_factor", &IterSplitExprNode::lower_factor)
        .def_ro("extent", &IterSplitExprNode::extent)
        .def_ro("scale", &IterSplitExprNode::scale);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSplitExpr", IterSplitExprNode, IterMapExprNode);
};
 
class IterSplitExpr : public IterMapExpr {
 public:
  TVM_DLL explicit IterSplitExpr(IterMark source);
  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr scale);
  TVM_DLL explicit IterSplitExpr(IterMark source, PrimExpr lower_factor, PrimExpr extent,
                                 PrimExpr scale);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSplitExpr, IterMapExpr, IterSplitExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode);
};
 
class IterSumExprNode : public IterMapExprNode {
 public:
  ffi::Array<IterSplitExpr> args;
  PrimExpr base;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IterSumExprNode>()
        .def_ro("args", &IterSumExprNode::args)
        .def_ro("base", &IterSumExprNode::base);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterSumExpr", IterSumExprNode, IterMapExprNode);
};
 
class IterSumExpr : public IterMapExpr {
 public:
  TVM_DLL IterSumExpr(ffi::Array<IterSplitExpr> args, PrimExpr base);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IterSumExpr, IterMapExpr, IterSumExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode);
};
 
}  // namespace arith
 
namespace ffi {
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSplitExprNode> = true;
template <>
inline constexpr bool object_ref_contains_v<PrimExpr, arith::IterSumExprNode> = true;
}  // namespace ffi
 
namespace arith {
 
enum IterMapLevel {
  // Require the mapping to be bijective.
  Bijective = 0,
  // Require the mapping to be surjective.
  Surjective = 1,
  // No mapping safety check.
  NoCheck = 3
};
 
class IterMapResultNode : public ffi::Object {
 public:
  // The detected pattern if a match exists.
  ffi::Array<IterSumExpr> indices;
 
  // Any errors that occurred while converting the input indices.  If
  // the array is empty, the conversion was successful.
  ffi::Array<ffi::String> errors;
 
  PrimExpr padding_predicate;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IterMapResultNode>()
        .def_ro("indices", &IterMapResultNode::indices)
        .def_ro("errors", &IterMapResultNode::errors)
        .def_ro("padding_predicate", &IterMapResultNode::padding_predicate);
  }
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("arith.IterMapResult", IterMapResultNode, ffi::Object);
};
 
class IterMapResult : public ffi::ObjectRef {
 public:
  // constructor
  IterMapResult() { data_ = ffi::make_object<IterMapResultNode>(); }
 
  IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
};
 
IterMapResult DetectIterMap(const ffi::Array<PrimExpr>& indices,
                            const ffi::Map<Var, Range>& input_iters, const PrimExpr& predicate,
                            IterMapLevel check_level, const arith::Analyzer& analyzer,
                            bool simplify_trivial_iterators = true);
 
ffi::Array<PrimExpr> IterMapSimplify(const ffi::Array<PrimExpr>& indices,
                                     const ffi::Map<Var, Range>& input_iters,
                                     const PrimExpr& input_pred, IterMapLevel check_level,
                                     const arith::Analyzer& analyzer,
                                     bool simplify_trivial_iterators = true);
 
ffi::Map<Var, PrimExpr> InverseAffineIterMap(const ffi::Array<IterSumExpr>& iter_map,
                                             const ffi::Array<PrimExpr> outputs);
 
ffi::Array<ffi::Array<IterMark>> SubspaceDivide(const ffi::Array<PrimExpr>& bindings,
                                                const ffi::Map<Var, Range>& input_iters,
                                                const ffi::Array<Var>& sub_iters,
                                                const PrimExpr& predicate, IterMapLevel check_level,
                                                const arith::Analyzer& analyzer,
                                                bool simplify_trivial_iterators = true);
 
PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr);
 
IterSumExpr NormalizeToIterSum(PrimExpr index, const ffi::Map<Var, Range>& input_iters,
                               const arith::Analyzer& analyzer);
 
}  // namespace arith
}  // namespace tvm
#endif  // TVM_ARITH_ITER_AFFINE_MAP_H_