block_scope.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_TIR_SCHEDULE_BLOCK_SCOPE_H_
#define TVM_TIR_SCHEDULE_BLOCK_SCOPE_H_

#include <tvm/tir/stmt.h>

#include <unordered_map>

namespace tvm {
namespace tir {

class StmtSRefNode : public Object {
 public:
  const StmtNode* stmt;
  StmtSRefNode* parent;
  int64_t seq_index;

  void VisitAttrs(AttrVisitor* v) {
    // `stmt` is not visited
    // `parent` is not visited
    v->Visit("seq_index", &seq_index);
  }

  static constexpr const char* _type_key = "tir.StmtSRef";
  TVM_DECLARE_FINAL_OBJECT_INFO(StmtSRefNode, Object);

  void Reset() {
    this->stmt = nullptr;
    this->parent = nullptr;
    this->seq_index = -1;
  }

  template <typename StmtType>
  const StmtType* StmtAs() const {
    if (stmt != nullptr && stmt->IsInstance<StmtType>()) {
      return static_cast<const StmtType*>(stmt);
    } else {
      return nullptr;
    }
  }
};

class StmtSRef : public ObjectRef {
 public:
  TVM_DLL explicit StmtSRef(const StmtNode* stmt, StmtSRefNode* parent, int64_t seq_index);

  StmtSRefNode* get() const { return static_cast<StmtSRefNode*>(data_.get()); }

  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(StmtSRef, ObjectRef, StmtSRefNode);

 public:
  TVM_DLL static StmtSRef InlineMark();
  TVM_DLL static StmtSRef RootMark();
};

enum class DepKind : int32_t {
  kRAW = 0,
  kWAW = 1,
  kWAR = 2,
  kOpaque = 3,
};

class DependencyNode : public Object {
 public:
  StmtSRef src;
  StmtSRef dst;
  DepKind kind;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("src", &src);
    v->Visit("dst", &dst);
    v->Visit("kind", &kind);
  }

  static constexpr const char* _type_key = "tir.Dependency";
  TVM_DECLARE_FINAL_OBJECT_INFO(DependencyNode, Object);
};

class Dependency : public ObjectRef {
 public:
  TVM_DLL explicit Dependency(StmtSRef src, StmtSRef dst, DepKind kind);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Dependency, ObjectRef, DependencyNode);
};

class BlockScopeNode : public Object {
 public:
  std::unordered_map<StmtSRef, Array<Dependency>, ObjectPtrHash, ObjectPtrEqual> src2deps;
  std::unordered_map<StmtSRef, Array<Dependency>, ObjectPtrHash, ObjectPtrEqual> dst2deps;
  std::unordered_map<Buffer, Array<StmtSRef>, ObjectPtrHash, ObjectPtrEqual> buffer_writers;
  bool stage_pipeline{false};

  void VisitAttrs(AttrVisitor* v) {}

  static constexpr const char* _type_key = "tir.BlockScope";
  TVM_DECLARE_FINAL_OBJECT_INFO(BlockScopeNode, Object);

 public:
  /******** Dependency ********/
  TVM_DLL Array<Dependency> GetDepsBySrc(const StmtSRef& src) const;
  TVM_DLL Array<Dependency> GetDepsByDst(const StmtSRef& dst) const;
};

class BlockScope : public ObjectRef {
 public:
  TVM_DLL BlockScope();
  TVM_DLL explicit BlockScope(const Array<StmtSRef>& child_block_srefs);

  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(BlockScope, ObjectRef, BlockScopeNode);
};

}  // namespace tir
}  // namespace tvm

#endif  // TVM_TIR_SCHEDULE_BLOCK_SCOPE_H_