tvm::tir::usmp Namespace Reference

Namespaces
	algo
	transform

Classes
class	AllocatedPoolInfo
struct	AllocatedPoolInfoNode
	This object contains information post-allocation for PoolInfo objects. More...
class	BufferInfo
class	BufferInfoAnalysis
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...
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	PoolAllocation
struct	PoolAllocationNode
	The pool allocation produced after the USMP algorithm. More...

Enumerations
enum	BufferInfoKind { BufferInfoKind::kIntermediate = 0, BufferInfoKind::kInput = 1, BufferInfoKind::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 AllocateNode *

op

)

Calculate the size of the extents in bytes.

Parameters

op	the allocate node

CalculateExtentsSize() [2/2]

Integer tvm::tir::usmp::CalculateExtentsSize

(

const AllocateConstNode *

op

)

Calculate the size of the extents in bytes.

Parameters

op	the allocate const 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