tvm::tir::usmp Namespace Reference

Namespaces
	algo
	transform

Classes
struct	BufferInfoNode
	Describes an abstract memory buffer that will get allocated inside a pool. The actual memory buffer in represented by PoolAllocationNode after static memory planning. More...
class	BufferInfo
struct	BufferInfoAnalysisNode
	This is a composite node that is produced by extract_buffer_info analysis pass that contains useful global information that could be useful for memory planning algorithms. More...
class	BufferInfoAnalysis
struct	PoolAllocationNode
	The pool allocation produced after the USMP algorithm. More...
class	PoolAllocation
struct	AllocatedPoolInfoNode
	This object contains information post-allocation for PoolInfo objects. More...
class	AllocatedPoolInfo

Enumerations
enum class	BufferInfoKind { kIntermediate = 0 , kInput = 1 , kOutput = 2 }
	A special kind to distinguish between I/O tensors to the model and intermediate tensors of the model. More...

Functions
BufferInfoAnalysis	ExtractBufferInfo (const PrimFunc &main_func, const IRModule &mod)
	Extract BufferInfo objects from a TIR IRModule. More...
Array< BufferInfo >	ConvertToArrayOfBufferInfo (const Map< BufferInfo, Stmt > &buffer_info_map)
	Convert the IR-bound BufferInfo map to an array of BufferInfo. More...
Integer	CalculateModuleWorkspaceSize (const IRModule &mod)
	Calculate workspace required to execute a IRModule with main expressed in TIR. More...
Integer	CalculateExtentsSize (const AllocateNode *op)
	Calculate the size of the extents in bytes. More...
Integer	CalculateExtentsSize (const AllocateConstNode *op)
	Calculate the size of the extents in bytes. More...
Map< Stmt, PoolAllocation >	AssignStmtPoolAllocations (const Map< BufferInfo, Stmt > &buffer_info_to_stmt, const Map< BufferInfo, PoolAllocation > &buffer_info_to_pool_allocation)
	Joins the Stmt nodes with PoolAllocation objects. More...
Map< String, PoolAllocation >	GetIOPoolAllocations (const Map< BufferInfo, PoolAllocation > &buffer_info_to_pool_allocation)
	Obtains I/O tensor names to their PoolAllocation objects. More...

Enumeration Type Documentation

BufferInfoKind

enum tvm::tir::usmp::BufferInfoKind

strong

A special kind to distinguish between I/O tensors to the model and intermediate tensors of the model.

Enumerator
kIntermediate
kInput
kOutput

Function Documentation

AssignStmtPoolAllocations()

Map<Stmt, PoolAllocation> tvm::tir::usmp::AssignStmtPoolAllocations	(	const Map< BufferInfo, Stmt > &	buffer_info_to_stmt,
		const Map< BufferInfo, PoolAllocation > &	buffer_info_to_pool_allocation
	)

Joins the Stmt nodes with PoolAllocation objects.

Parameters

buffer_info_to_stmt	the map of BufferInfo objects to Stmt nodes
buffer_info_to_pool_allocation	the map of BufferInfo objects to PoolAllocation objects

CalculateExtentsSize() [1/2]

Integer tvm::tir::usmp::CalculateExtentsSize

(

const AllocateConstNode *

op

)

Calculate the size of the extents in bytes.

Parameters

op	the allocate const node

CalculateExtentsSize() [2/2]

Integer tvm::tir::usmp::CalculateExtentsSize

(

const AllocateNode *

op

)

Calculate the size of the extents in bytes.

Parameters

op	the allocate node

CalculateModuleWorkspaceSize()

Integer tvm::tir::usmp::CalculateModuleWorkspaceSize

(

const IRModule &

mod

)

Calculate workspace required to execute a IRModule with main expressed in TIR.

Parameters

mod	the IRModule with TIR-based main function

ConvertToArrayOfBufferInfo()

Array<BufferInfo> tvm::tir::usmp::ConvertToArrayOfBufferInfo

(

const Map< BufferInfo, Stmt > &

buffer_info_map

)

Convert the IR-bound BufferInfo map to an array of BufferInfo.

Parameters

buffer_info_map

IR-bound BufferInfo map

ExtractBufferInfo()

BufferInfoAnalysis tvm::tir::usmp::ExtractBufferInfo	(	const PrimFunc &	main_func,
		const IRModule &	mod
	)

Extract BufferInfo objects from a TIR IRModule.

This pass would extract the buffer information of allocate nodes including liveness conflict with other buffer info objects.

Returns

A Map of BufferInfo objects and their associated Stmts

GetIOPoolAllocations()

Map<String, PoolAllocation> tvm::tir::usmp::GetIOPoolAllocations

(

const Map< BufferInfo, PoolAllocation > &

buffer_info_to_pool_allocation

)

Obtains I/O tensor names to their PoolAllocation objects.

Parameters

buffer_info_to_pool_allocation

the map of BufferInfo objects to PoolAllocation objects

This function will obtain pool allocations for I/O tensors if that had been planned