- Feature Name: Adding annotatation field to tir.allocate nodes
- Start Date: 2021-06-01
- RFC PR: #23
- GitHub Issue: TBD
This RFC proposes to annotation field tir.allocate nodes. These annotations can be used as auxiliary hint to future transformations.
Currently, TVM relies on dynamic (alloc and free style) allocations in runtime to manage the intermediary memory used by operators and the network. This is sometimes not desirable, especially in microTVM.
The current design of Unified Static Memory Planner (USMP), enables the user option to provide buffers to place workspace and constant tensors.
tvmc compile my_model.tflite
--executor=aot
--target=accel,c
--with-workspace-buffer= "name=dtcm;target=c;size=1000" # Here the size is more of a hint/guide provided to USMP
--with-workspace-buffer= "name=sram;target=c,accel"
--with-parameter-buffer= "name=itcm;target=c;size=5000" # Here the size is more of a hint/guide provided to USMP
--with-parameter-buffer= "name=flash;target=c,accel"
// The User Application
extern const TVMModel my_model;
__attribute__((section( "ITCM" ) const uint8_t my_model_params_1[TVM_MY_MODEL_ITCM_PARAMETER_BUFFER_SIZE] = <param_1_data>;
__attribute__((section( "FLASH" ), aligned( 16 ))) const uint8_t my_model_params_2[TVM_MY_MODEL_FLASH_PARAMETER_BUFFER_SIZE] = <param_2_data>;
__attribute__((section( "DTCM" ) static uint8_t workspace_buffer_1[TVM_MY_MODEL_DTCM_WORKSPACE_BUFFER_SIZE];
__attribute__((section( "SRAM" ), aligned( 16 ))) static uint8_t workspace_buffer_2[TVM_MY_MODEL_SRAM_WORKSPACE_BUFFER_SIZE];
int main(...) {
...
TVMContext context;
TVMInputs_my_model_1 inputs = {input};
TVMOutputs_my_model_1 outputs = {output};
TVMWorkspaces_my_model workspaces = {
.sram = &workspace_buffer_1,
.dtcm = &workspace_buffer_2,
};
TVMParameters_my_model parameters = {
.flash = &my_model_params_1,
.itcm = &my_model_params_2
};
TVMSetWorkspaces(&context, &workspaces);
TVMSetParameters(&context, parameters);
TVMExecute(&my_model, &inputs, &outputs, &context);
}
Therefore, we'd need a way to represent the association of each of these memories, that the user will pin the buffers to, closer to allocate nodes in TIR.
At the IR, we ll need to associate each allocate node with one (or more) memories that it can end up, because the scheduling might be satisfied with placing buffers in any of the memories in a given set of memories. Therefore, the scheduler might want the memory planners to decide which memory to use based on finding the allocation that fit.
There are broadly two requirements here :
P1 : Indicate candidate memories (a.k.a. Pools) that each allocate be associated with
P2 : Indicate the final memory the allocate will be pinned on
To serve P2, we propose to use the existing tag of the storage_scope with 'global-<memory_name>'.
struct StorageScope {
/*! \brief The rank of the storage */
StorageRank rank{StorageRank::kGlobal};
/*! \brief tag for special purpose memory. */
std::string tag;
To serve P1, this RFC introduces the addition of the annotations field
This is not particularly a user-facing feature.
To serve P1, we propose to use :
class AllocateNode : public StmtNode {
public:
/*! \brief The buffer variable. */
Var buffer_var;
/*! \brief The type of the buffer. */
DataType dtype;
/*! \brief The extents of the buffer. */
Array<PrimExpr> extents;
/*! \brief Only allocate buffer when condition is satisfied. */
PrimExpr condition;
/*! \brief The body to be executed. */
Stmt body;
/*! \brief If the allocate is scoped global, this field indicates
* which external memories it could be pinned to as a comma seperated
* string.
*/
+ Map<String, Objectref> annotations;
Here the addition of the annotations field could serve offer hints/guides to future passes (i.e. In Unified Static Memory Planner, we could use "candidate_memory_pools" as the key while the value being Map<String, PoolInfo>.)
TIR:
allocate_node_1 = tir.allocate([157323], "int16", "global")
tir.attr(allocate_node_1, "pinned_memory", "foo_memory,bar_memory")
class AllocateNode : public StmtNode {
public:
/*! \brief The buffer variable. */
Var buffer_var;
/*! \brief The type of the buffer. */
DataType dtype;
/*! \brief The extents of the buffer. */
Array<PrimExpr> extents;
/*! \brief Only allocate buffer when condition is satisfied. */
PrimExpr condition;
/*! \brief The body to be executed. */
Stmt body;
/*! \brief If the allocate is scoped global, this field indicates
* which external memories it could be pinned to as a comma seperated
* string.
*/
String pinned_memory;
TIR:
allocate_node_1 = tir.allocate([157323], "int16", "global", "foo_memory,bar_memory")
/*! \brief class to represent storage scope */
struct StorageScope {
/*! \brief The rank of the storage */
StorageRank rank{StorageRank::kGlobal};
/*! \brief tag for special purpose memories. */
Array<String> tags;
/*!
* \brief Create storage scope from string
* \param s The string to be parsed.
* \return The storage scope.
*/
static StorageScope Create(const std::string& s) {
StorageScope r;
if (s.empty()) {/
r.rank = StorageRank::kGlobal;
} else if (s.compare(0, 6, "global") == 0) {
r.rank = StorageRank::kGlobal;
r.tags = parseTags(s);
TIR:
allocate_node_1 = tir.allocate([157323], "int16", "global.(foo_memory,bar_memory)")
Out of the options, S1 seems the most non-invasive. However, we will need special handlers to obtain the information.
S2 does not change the storage scope (or the 'tags') and adds an additional field to allocates to note this information.
S3 fold the information into storage_scope and utilizes the 'tag' denote the memory. The change to IR, is just to support more tags.
However, in the discussion with community, we felt the need to seperate changes that serves P1 and P2, respectively. In fact, for P2 we could re-use the tag of storage_scope without an IR change. For P1, it seems a good and general change to include annotations for the allocate node.
None. Its consistent with rest of the IR design and allows other features to use to pass hints for future transformation in the compiler.