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/ir/diagnostic.h>
#include <tvm/ir/expr.h>
#include <tvm/tir/var.h>
 
namespace tvm {
namespace arith {
 
class IterMapExprNode : public PrimExprNode {
 public:
  // overrides
  void VisitAttrs(tvm::AttrVisitor* v) {}
 
  static constexpr const char* _type_key = "arith.IterMapExpr";
  static constexpr const uint32_t _type_child_slots = 3;
  TVM_DECLARE_BASE_OBJECT_INFO(IterMapExprNode, PrimExprNode);
};
 
class IterMapExpr : public PrimExpr {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(IterMapExpr, PrimExpr, IterMapExprNode);
};
 
class IterMarkNode : public Object {
 public:
  PrimExpr source;
  PrimExpr extent;
 
  // overrides
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("source", &source);
    v->Visit("extent", &extent);
  }
 
  bool SEqualReduce(const IterMarkNode* other, SEqualReducer equal) const {
    equal->MarkGraphNode();
    return equal(source, other->source) && equal(extent, other->extent);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce->MarkGraphNode();
    hash_reduce(source);
    hash_reduce(extent);
  }
 
  static constexpr const bool _type_has_method_sequal_reduce = true;
  static constexpr const bool _type_has_method_shash_reduce = true;
  static constexpr const char* _type_key = "arith.IterMark";
  TVM_DECLARE_FINAL_OBJECT_INFO(IterMarkNode, Object);
};
 
class IterMark : public ObjectRef {
 public:
  TVM_DLL IterMark(PrimExpr source, PrimExpr extent);
 
  TVM_DEFINE_OBJECT_REF_METHODS(IterMark, ObjectRef, IterMarkNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterMarkNode);
};
 
class IterSplitExprNode : public IterMapExprNode {
 public:
  IterMark source;
  PrimExpr lower_factor;
  PrimExpr extent;
  PrimExpr scale;
 
  // overrides
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("source", &source);
    v->Visit("lower_factor", &lower_factor);
    v->Visit("extent", &extent);
    v->Visit("scale", &scale);
  }
 
  bool SEqualReduce(const IterSplitExprNode* other, SEqualReducer equal) const {
    return equal(source, other->source) && equal(lower_factor, other->lower_factor) &&
           equal(extent, other->extent) && equal(scale, other->scale);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(source);
    hash_reduce(lower_factor);
    hash_reduce(extent);
    hash_reduce(scale);
  }
 
  static constexpr const char* _type_key = "arith.IterSplitExpr";
  TVM_DECLARE_FINAL_OBJECT_INFO(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_DEFINE_OBJECT_REF_METHODS(IterSplitExpr, IterMapExpr, IterSplitExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSplitExprNode);
};
 
class IterSumExprNode : public IterMapExprNode {
 public:
  Array<IterSplitExpr> args;
  PrimExpr base;
 
  // overrides
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("args", &args);
    v->Visit("base", &base);
  }
 
  bool SEqualReduce(const IterSumExprNode* other, SEqualReducer equal) const {
    return equal(args, other->args) && equal(base, other->base);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(args);
    hash_reduce(base);
  }
 
  static constexpr const char* _type_key = "arith.IterSumExpr";
  TVM_DECLARE_FINAL_OBJECT_INFO(IterSumExprNode, IterMapExprNode);
};
 
class IterSumExpr : public IterMapExpr {
 public:
  TVM_DLL IterSumExpr(Array<IterSplitExpr> args, PrimExpr base);
 
  TVM_DEFINE_OBJECT_REF_METHODS(IterSumExpr, IterMapExpr, IterSumExprNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(IterSumExprNode);
};
 
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 Object {
 public:
  // The detected pattern if a match exists.
  Array<IterSumExpr> indices;
 
  // Any errors that occurred while converting the input indices.  If
  // the array is empty, the conversion was successful.
  Array<String> errors;
 
  PrimExpr padding_predicate;
 
  // overrides
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("errors", &errors);
    v->Visit("indices", &indices);
    v->Visit("padding_predicate", &padding_predicate);
  }
 
  static constexpr const char* _type_key = "arith.IterMapResult";
  TVM_DECLARE_FINAL_OBJECT_INFO(IterMapResultNode, Object);
};
 
class IterMapResult : public ObjectRef {
 public:
  // constructor
  IterMapResult() { data_ = make_object<IterMapResultNode>(); }
 
  IterMapResultNode* operator->() const { return static_cast<IterMapResultNode*>(get_mutable()); }
};
 
IterMapResult DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
                            const PrimExpr& predicate, IterMapLevel check_level,
                            arith::Analyzer* analyzer, bool simplify_trivial_iterators = true);
 
Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
                                const PrimExpr& input_pred, IterMapLevel check_level,
                                arith::Analyzer* analyzer, bool simplify_trivial_iterators = true);
 
Map<Var, PrimExpr> InverseAffineIterMap(const Array<IterSumExpr>& iter_map,
                                        const Array<PrimExpr> outputs);
 
Array<Array<IterMark>> SubspaceDivide(const Array<PrimExpr>& bindings,
                                      const Map<Var, Range>& input_iters,
                                      const Array<Var>& sub_iters, const PrimExpr& predicate,
                                      IterMapLevel check_level, arith::Analyzer* analyzer,
                                      bool simplify_trivial_iterators = true);
 
PrimExpr NormalizeIterMapToExpr(const PrimExpr& expr);
 
IterSumExpr NormalizeToIterSum(PrimExpr index, const Map<Var, Range>& input_iters,
                               arith::Analyzer* analyzer);
 
}  // namespace arith
}  // namespace tvm
#endif  // TVM_ARITH_ITER_AFFINE_MAP_H_