utils.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_USMP_UTILS_H_
#define TVM_TIR_USMP_UTILS_H_
 
#include <tvm/ir/expr.h>
#include <tvm/ir/memory_pools.h>
#include <tvm/runtime/device_api.h>
#include <tvm/target/target.h>
#include <tvm/tir/stmt.h>
 
namespace tvm {
 
constexpr const char* kUSMPEnableOption = "tir.usmp.enable";
constexpr const char* kUSMPAlgorithmOption = "tir.usmp.algorithm";
constexpr const char* kUSMPUseWorkspaceIO = "tir.usmp.use_workspace_io";
constexpr const char* kUSMPCustomAlgorithmOption = "tir.usmp.custom_algorithm";
 
namespace tir {
namespace usmp {
enum class BufferInfoKind { kIntermediate = 0, kInput = 1, kOutput = 2 };
 
struct BufferInfoNode : public Object {
  String name_hint;
  Integer size_bytes;
  Array<PoolInfo> pool_candidates;
  Integer alignment;
  Array<ObjectRef> conflicts;
  BufferInfoKind kind;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("name_hint", &name_hint);
    v->Visit("size_bytes", &size_bytes);
    v->Visit("pool_candidates", &pool_candidates);
    v->Visit("alignment", &alignment);
    v->Visit("conflicts", &conflicts);
    v->Visit("kind", &kind);
  }
 
  bool SEqualReduce(const BufferInfoNode* other, SEqualReducer equal) const {
    return equal(name_hint, other->name_hint) && equal(size_bytes, other->size_bytes) &&
           equal(pool_candidates, other->pool_candidates) && equal(alignment, other->alignment) &&
           equal(conflicts, other->conflicts) && equal(kind, other->kind);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(name_hint);
    hash_reduce(size_bytes);
    hash_reduce(alignment);
    hash_reduce(conflicts);
    hash_reduce(pool_candidates);
    hash_reduce(kind);
  }
  TVM_DLL void SetConflicts(Array<ObjectRef> conflicting_buffer_info_objs);
 
  static constexpr const char* _type_key = "tir.usmp.BufferInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(BufferInfoNode, Object);
};
 
class BufferInfo : public ObjectRef {
 public:
  TVM_DLL BufferInfo(String name_hint, Integer size_bytes, Array<PoolInfo> pool_candidates,
                     Integer alignment = runtime::kDefaultWorkspaceAlignment,
                     BufferInfoKind kind = BufferInfoKind::kIntermediate);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(BufferInfo, ObjectRef, BufferInfoNode);
};
 
struct BufferInfoAnalysisNode : public Object {
  Map<BufferInfo, tir::Stmt> buffer_info_stmts;
  Integer memory_pressure;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("buffer_info_stmts", &buffer_info_stmts);
    v->Visit("memory_pressure", &memory_pressure);
  }
 
  bool SEqualReduce(const BufferInfoAnalysisNode* other, SEqualReducer equal) const {
    return equal(buffer_info_stmts, other->buffer_info_stmts) &&
           equal(memory_pressure, other->memory_pressure);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(buffer_info_stmts);
    hash_reduce(memory_pressure);
  }
};
 
class BufferInfoAnalysis : public ObjectRef {
 public:
  TVM_DLL BufferInfoAnalysis(Map<BufferInfo, tir::Stmt> buffer_info_stmts, Integer memory_pressure);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(BufferInfoAnalysis, ObjectRef, BufferInfoAnalysisNode);
};
 
struct PoolAllocationNode : public Object {
  PoolInfo pool_info;
  Integer byte_offset;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pool_info", &pool_info);
    v->Visit("byte_offset", &byte_offset);
  }
 
  bool SEqualReduce(const PoolAllocationNode* other, SEqualReducer equal) const {
    return equal(pool_info, other->pool_info) && equal(byte_offset, other->byte_offset);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(pool_info);
    hash_reduce(byte_offset);
  }
 
  static constexpr const char* _type_key = "tir.usmp.PoolAllocation";
  TVM_DECLARE_FINAL_OBJECT_INFO(PoolAllocationNode, Object);
};
 
class PoolAllocation : public ObjectRef {
 public:
  TVM_DLL PoolAllocation(PoolInfo pool_info, Integer byte_offset);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolAllocation, ObjectRef, PoolAllocationNode);
};
 
struct AllocatedPoolInfoNode : public Object {
  PoolInfo pool_info;
  Integer allocated_size;
  Optional<Integer> pool_var_idx;
 
  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pool_info", &pool_info);
    v->Visit("allocated_size", &allocated_size);
    v->Visit("pool_var_idx", &pool_var_idx);
  }
 
  bool SEqualReduce(const AllocatedPoolInfoNode* other, SEqualReducer equal) const {
    return equal(pool_info, other->pool_info) && equal(allocated_size, other->allocated_size) &&
           equal(pool_var_idx, other->pool_var_idx);
  }
 
  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(pool_info);
    hash_reduce(allocated_size);
    hash_reduce(pool_var_idx);
  }
 
  static constexpr const char* _type_key = "ir.AllocatedPoolInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(AllocatedPoolInfoNode, Object);
};
 
class AllocatedPoolInfo : public ObjectRef {
 public:
  TVM_DLL AllocatedPoolInfo(PoolInfo pool_info, Integer allocated_size,
                            Integer pool_var_idx = Integer());
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(AllocatedPoolInfo, ObjectRef, AllocatedPoolInfoNode);
};
 
Array<BufferInfo> ConvertToArrayOfBufferInfo(const Map<BufferInfo, Stmt>& buffer_info_map);
 
Integer CalculateModuleWorkspaceSize(const IRModule& mod);
 
static constexpr const char* kPoolCandidatesAllocateAttr = "candidate_memory_pools";
 
static constexpr const char* kInputTensorAllocate = "input_tensor";
 
static constexpr const char* kOutputTensorAllocate = "output_tensor";
 
Integer CalculateExtentsSize(const AllocateNode* op);
 
Integer CalculateExtentsSize(const AllocateConstNode* op);
 
Map<Stmt, PoolAllocation> AssignStmtPoolAllocations(
    const Map<BufferInfo, Stmt>& buffer_info_to_stmt,
    const Map<BufferInfo, PoolAllocation>& buffer_info_to_pool_allocation);
 
Map<String, PoolAllocation> GetIOPoolAllocations(
    const Map<BufferInfo, PoolAllocation>& buffer_info_to_pool_allocation);
 
}  // namespace usmp
}  // namespace tir
 
namespace attr {
static constexpr const char* kPoolArgs = "pool_args";
 
static constexpr const char* kIOTensorPoolAllocations = "io_tensor_pool_allocations";
 
}  // namespace attr
 
}  // namespace tvm
 
#endif  // TVM_TIR_USMP_UTILS_H_