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/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);
};

Array<IterSumExpr> DetectIterMap(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
                                 const PrimExpr& predicate, bool require_bijective,
                                 arith::Analyzer* analyzer);
Array<PrimExpr> IterMapSimplify(const Array<PrimExpr>& indices, const Map<Var, Range>& input_iters,
                                const PrimExpr& input_pred, bool require_bijective);

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,
                                      bool require_bijective, arith::Analyzer* analyzer);

}  // namespace arith
}  // namespace tvm
#endif  // TVM_ARITH_ITER_AFFINE_MAP_H_