memory_pools.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_IR_MEMORY_POOLS_H_
#define TVM_IR_MEMORY_POOLS_H_

#include <tvm/runtime/registry.h>
#include <tvm/target/target.h>

struct TVMConstantInfo;
namespace tvm {

struct PoolInfoNode : public Object {
 public:
  String pool_name;
  Integer size_hint_bytes;
  Integer clock_frequency_hz;
  Integer read_bandwidth_bytes_per_cycle;
  Integer write_bandwidth_bytes_per_cycle;
  Integer read_latency_cycles;
  Integer write_latency_cycles;
  Map<Target, Integer> target_burst_bytes;
  bool is_internal = false;

  Array<Target> targets;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("pool_name", &pool_name);
    v->Visit("size_hint_bytes", &size_hint_bytes);
    v->Visit("clock_frequency_hz", &clock_frequency_hz);
    v->Visit("read_bandwidth_bytes_per_cycle", &read_bandwidth_bytes_per_cycle);
    v->Visit("write_bandwidth_bytes_per_cycle", &write_bandwidth_bytes_per_cycle);
    v->Visit("read_latency_cycles", &read_latency_cycles);
    v->Visit("write_latency_cycles", &write_latency_cycles);
    v->Visit("target_burst_bytes", &target_burst_bytes);
    v->Visit("is_internal", &is_internal);
  }

  bool SEqualReduce(const PoolInfoNode* other, SEqualReducer equal) const {
    return equal(pool_name, other->pool_name) && equal(size_hint_bytes, other->size_hint_bytes) &&
           equal(clock_frequency_hz, other->clock_frequency_hz) &&
           equal(read_bandwidth_bytes_per_cycle, other->read_bandwidth_bytes_per_cycle) &&
           equal(write_bandwidth_bytes_per_cycle, other->write_bandwidth_bytes_per_cycle) &&
           equal(read_latency_cycles, other->read_latency_cycles) &&
           equal(write_latency_cycles, other->write_latency_cycles) &&
           equal(target_burst_bytes, other->target_burst_bytes) &&
           equal(is_internal, other->is_internal);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(pool_name);
    hash_reduce(size_hint_bytes);
    hash_reduce(clock_frequency_hz);
    hash_reduce(read_bandwidth_bytes_per_cycle);
    hash_reduce(write_bandwidth_bytes_per_cycle);
    hash_reduce(read_latency_cycles);
    hash_reduce(write_latency_cycles);
    hash_reduce(target_burst_bytes);
    hash_reduce(is_internal);
  }

  static constexpr const char* _type_key = "ir.PoolInfo";
  TVM_DECLARE_BASE_OBJECT_INFO(PoolInfoNode, Object);
};

static constexpr const char* kTargetPoolReadWriteAccess = "rw";

static constexpr const char* kTargetPoolReadOnlyAccess = "ro";

static const int kUnrestrictedPoolSizeHint = -1;

static const int kUnknownClockFrequency = -1;

static const int kUnknownReadBandwidth = -1;

static const int kUnknownWriteBandwidth = -1;

class PoolInfo : public ObjectRef {
 protected:
  TVM_DLL PoolInfo(String pool_name, Integer size_hint_bytes = kUnrestrictedPoolSizeHint,
                   Integer clock_frequency_hz = kUnknownClockFrequency,
                   Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth,
                   Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth,
                   Integer read_latency_cycles = 0, Integer write_latency_cycles = 0,
                   Map<Target, Integer> target_burst_bytes = {}, Bool is_internal = Bool(false));

 public:
  TVM_DEFINE_OBJECT_REF_METHODS(PoolInfo, ObjectRef, PoolInfoNode);
};

struct PoolInfoPropertiesNode : public Object {
  Integer size_hint_bytes = kUnrestrictedPoolSizeHint;
  Integer clock_frequency_hz = kUnknownClockFrequency;
  Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth;
  Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth;
  Integer read_latency_cycles = 0;
  Integer write_latency_cycles = 0;
  Map<Target, Integer> target_burst_bytes{};
  bool is_internal = false;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("size_hint_bytes", &size_hint_bytes);
    v->Visit("clock_frequency_hz", &clock_frequency_hz);
    v->Visit("read_bandwidth_bytes_per_cycle", &read_bandwidth_bytes_per_cycle);
    v->Visit("write_bandwidth_bytes_per_cycle", &write_bandwidth_bytes_per_cycle);
    v->Visit("read_latency_cycles", &read_latency_cycles);
    v->Visit("write_latency_cycles", &write_latency_cycles);
    v->Visit("target_burst_bytes", &target_burst_bytes);
    v->Visit("is_internal", &is_internal);
  }

  bool SEqualReduce(const PoolInfoPropertiesNode* other, SEqualReducer equal) const {
    return equal(size_hint_bytes, other->size_hint_bytes) &&
           equal(clock_frequency_hz, other->clock_frequency_hz) &&
           equal(read_bandwidth_bytes_per_cycle, other->read_bandwidth_bytes_per_cycle) &&
           equal(write_bandwidth_bytes_per_cycle, other->write_bandwidth_bytes_per_cycle) &&
           equal(read_latency_cycles, other->read_latency_cycles) &&
           equal(write_latency_cycles, other->write_latency_cycles) &&
           equal(target_burst_bytes, other->target_burst_bytes) &&
           equal(is_internal, other->is_internal);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(size_hint_bytes);
    hash_reduce(clock_frequency_hz);
    hash_reduce(read_bandwidth_bytes_per_cycle);
    hash_reduce(write_bandwidth_bytes_per_cycle);
    hash_reduce(read_latency_cycles);
    hash_reduce(write_latency_cycles);
    hash_reduce(target_burst_bytes);
    hash_reduce(is_internal);
  }

  static constexpr const char* _type_key = "ir.PoolInfoProperties";
  TVM_DECLARE_FINAL_OBJECT_INFO(PoolInfoPropertiesNode, Object);
};

class PoolInfoProperties : public ObjectRef {
 public:
  TVM_DLL PoolInfoProperties(Integer size_hint_bytes,
                             Integer clock_frequency_hz = kUnknownClockFrequency,
                             Integer read_bandwidth_bytes_per_cycle = kUnknownReadBandwidth,
                             Integer write_bandwidth_bytes_per_cycle = kUnknownWriteBandwidth,
                             Integer read_latency_cycles = 0, Integer write_latency_cycles = 0,
                             Map<Target, Integer> target_burst_bytes = {},
                             Bool is_internal = Bool(false));
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(PoolInfoProperties, ObjectRef, PoolInfoPropertiesNode);
};

/* \brief Represents RW memory area */
struct WorkspacePoolInfoNode : public PoolInfoNode {
  void VisitAttrs(tvm::AttrVisitor* v) { PoolInfoNode::VisitAttrs(v); }

  bool SEqualReduce(const WorkspacePoolInfoNode* other, SEqualReducer equal) const {
    return PoolInfoNode::SEqualReduce(other, equal);
  }

  void SHashReduce(SHashReducer hash_reduce) const { PoolInfoNode::SHashReduce(hash_reduce); }

  static constexpr const char* _type_key = "ir.WorkspacePoolInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(WorkspacePoolInfoNode, PoolInfoNode);
};

class WorkspacePoolInfo : public PoolInfo {
 public:
  TVM_DLL WorkspacePoolInfo(
      String pool_name, Array<Target> targets,
      PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint));
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(WorkspacePoolInfo, PoolInfo, WorkspacePoolInfoNode);
};

/*
 * \brief The ConstantInfoNode contains numeric literal in RO pool
 * Used to initialise RO memory in ConstantPoolInfo
 */
struct ConstantInfoNode : public Object {
  String name_hint;
  Integer byte_offset;
  runtime::NDArray data;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("name_hint", &name_hint);
    v->Visit("byte_offset", &byte_offset);
    v->Visit("data", &data);
  }

  bool SEqualReduce(const ConstantInfoNode* other, SEqualReducer equal) const {
    return equal(name_hint, other->name_hint) && equal(byte_offset, other->byte_offset) &&
           equal(data, other->data);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce(name_hint);
    hash_reduce(byte_offset);
    hash_reduce(data);
  }

  static constexpr const char* _type_key = "ir.ConstantInfo";
  static constexpr bool _type_has_method_sequal_reduce = true;
  static constexpr bool _type_has_method_shash_reduce = true;
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantInfoNode, Object);
};

class ConstantInfo : public ObjectRef {
 public:
  TVM_DLL ConstantInfo(const struct ::TVMConstantInfo* data);
  ConstantInfo(String name, Integer byte_offset, runtime::NDArray data);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantInfo, ObjectRef, ConstantInfoNode);
};

/* \brief ConstantPoolInfoNode represents an RO memory area initialized with
 * data from constant_info_array */
struct ConstantPoolInfoNode : public PoolInfoNode {
  Array<ConstantInfo> constant_info_array;

  void VisitAttrs(tvm::AttrVisitor* v) {
    PoolInfoNode::VisitAttrs(v);
    v->Visit("constant_info_array", &constant_info_array);
  }

  bool SEqualReduce(const ConstantPoolInfoNode* other, SEqualReducer equal) const {
    return PoolInfoNode::SEqualReduce(other, equal) &&
           equal(constant_info_array, other->constant_info_array);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    PoolInfoNode::SHashReduce(hash_reduce);
    hash_reduce(constant_info_array);
  }

  static constexpr const char* _type_key = "ir.ConstantPoolInfo";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPoolInfoNode, PoolInfoNode);
};

class ConstantPoolInfo : public PoolInfo {
 public:
  TVM_DLL ConstantPoolInfo(
      String pool_name, Array<Target> targets, Array<ConstantInfo> constant_info_array,
      PoolInfoProperties properties = PoolInfoProperties(kUnrestrictedPoolSizeHint));
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantPoolInfo, PoolInfo, ConstantPoolInfoNode);
};

/* \brief A container for WorkspacePoolInfo objects */
struct WorkspaceMemoryPoolsNode : public Object {
  Array<PoolInfo> pools;

  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("pools", &pools); }

  bool SEqualReduce(const WorkspaceMemoryPoolsNode* other, SEqualReducer equal) const {
    return equal(pools, other->pools);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(pools); }

  static constexpr const char* _type_key = "ir.WorkspaceMemoryPools";
  TVM_DECLARE_FINAL_OBJECT_INFO(WorkspaceMemoryPoolsNode, Object);
};

class WorkspaceMemoryPools : public ObjectRef {
 public:
  TVM_DLL WorkspaceMemoryPools(Array<PoolInfo> pools);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(WorkspaceMemoryPools, ObjectRef, WorkspaceMemoryPoolsNode);
};

/* \brief A container for ConstantPoolInfo objects */
struct ConstantMemoryPoolsNode : public Object {
  Array<ConstantPoolInfo> pools;

  void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("pools", &pools); }

  bool SEqualReduce(const ConstantMemoryPoolsNode* other, SEqualReducer equal) const {
    return equal(pools, other->pools);
  }

  void SHashReduce(SHashReducer hash_reduce) const { hash_reduce(pools); }

  static constexpr const char* _type_key = "ir.ConstantMemoryPools";
  TVM_DECLARE_FINAL_OBJECT_INFO(ConstantMemoryPoolsNode, Object);
};

class ConstantMemoryPools : public ObjectRef {
 public:
  TVM_DLL ConstantMemoryPools(Array<ConstantPoolInfo> pools);
  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(ConstantMemoryPools, ObjectRef, ConstantMemoryPoolsNode);
};

}  // namespace tvm

#endif  // TVM_IR_MEMORY_POOLS_H_