tvm
Functions | Variables
tvm::tirx::attr Namespace Reference

PrimFunc specific attribute names. More...

Functions

bool IsPragmaKey (const std::string &attr_key)
 Check if attr_key is a pragma key extension. More...
 

Variables

constexpr const char * kKernelLaunchParams = "tirx.kernel_launch_params"
 List of thread IterVar that a DeviceLaunch function corresponds to. More...
 
constexpr const char * kNoAlias = "tirx.noalias"
 Whether to set noalias rule on the function arguments. More...
 
constexpr const char * kIsEntryFunc = "tirx.is_entry_func"
 Mark the function as the entry function of the final generated runtime module. More...
 
constexpr const char * kIsGlobalFunc = "tirx.is_global_func"
 Mark the function as the global function called from the host. More...
 
constexpr const char * kIsHostFunc = "tirx.is_host_func"
 Mark the function as run on the host, mutually exclusive with kTarget. More...
 
constexpr const char * kIsScheduled = "tirx.is_scheduled"
 Mark the function as scheduled, so the default schedule will pass will skip it. More...
 
constexpr const char * buffer_bound = "buffer_bound"
 Mark stores/loads with their bounds. More...
 
constexpr const char * compute_scope = "compute_scope"
 Mark the scope as when computation start to happen. This can hint some code generator to create a new function for compute. More...
 
constexpr const char * device_id = "device_id"
 The allocation device for global malloc in host. More...
 
constexpr const char * device_scope = "device_scope"
 Mark that it is in the device scope. More...
 
constexpr const char * device_type = "device_type"
 The device type. More...
 
constexpr const char * extern_scope = "extern_scope"
 Mark the scope as generated by extern primitive. Such scope can contain arbitrary ir program and we need to be careful when making certain assumptions about the structure of the program. More...
 
constexpr const char * pragma_auto_unroll_max_step = "pragma_auto_unroll_max_step"
 Pragma: auto-unroll, max_step. More...
 
constexpr const char * pragma_import_c = "pragma_import_c"
 Import C source or file into the final code gen module. More...
 
constexpr const char * pragma_import_llvm = "pragma_import_llvm"
 Import llvm source or file into the final code gen module. More...
 
constexpr const char * pragma_scope_prefix = "pragma_"
 Mark region is guarded by the pragma extension. More...
 
constexpr const char * pragma_tensor_core = "pragma_tensor_core"
 Try to modify the AST to support Tensor Core. More...
 
constexpr const char * pragma_unroll_explicit = "pragma_unroll_explicit"
 Pragma: unroll explicit. More...
 
constexpr const char * storage_alignment = "storage_alignment"
 Mark storage alignment requirement of buffers. More...
 
constexpr const char * thread_extent = "thread_extent"
 Mark launching extent of thread, used by device API. More...
 
constexpr const char * kVolatile = "tirx.volatile"
 Annotation key on AllocBuffer marking the allocation as volatile. More...
 

Detailed Description

PrimFunc specific attribute names.

namespace of possible attributes in AttrStmt.attr_key

See also
tvm::attr

Function Documentation

◆ IsPragmaKey()

bool tvm::tirx::attr::IsPragmaKey ( const std::string &  attr_key)
inline

Check if attr_key is a pragma key extension.

Parameters
attr_keyThe attr key to be compared
Returns
true if it is a pragma key

Variable Documentation

◆ buffer_bound

constexpr const char* tvm::tirx::attr::buffer_bound = "buffer_bound"
constexpr

Mark stores/loads with their bounds.

◆ compute_scope

constexpr const char* tvm::tirx::attr::compute_scope = "compute_scope"
constexpr

Mark the scope as when computation start to happen. This can hint some code generator to create a new function for compute.

◆ device_id

constexpr const char* tvm::tirx::attr::device_id = "device_id"
constexpr

The allocation device for global malloc in host.

◆ device_scope

constexpr const char* tvm::tirx::attr::device_scope = "device_scope"
constexpr

Mark that it is in the device scope.

◆ device_type

constexpr const char* tvm::tirx::attr::device_type = "device_type"
constexpr

The device type.

◆ extern_scope

constexpr const char* tvm::tirx::attr::extern_scope = "extern_scope"
constexpr

Mark the scope as generated by extern primitive. Such scope can contain arbitrary ir program and we need to be careful when making certain assumptions about the structure of the program.

◆ kIsEntryFunc

constexpr const char* tvm::tirx::attr::kIsEntryFunc = "tirx.is_entry_func"
constexpr

Mark the function as the entry function of the final generated runtime module.

Type: Integer

Note
There can only be one entry function per module.

◆ kIsGlobalFunc

constexpr const char* tvm::tirx::attr::kIsGlobalFunc = "tirx.is_global_func"
constexpr

Mark the function as the global function called from the host.

Type: Integer

◆ kIsHostFunc

constexpr const char* tvm::tirx::attr::kIsHostFunc = "tirx.is_host_func"
constexpr

Mark the function as run on the host, mutually exclusive with kTarget.

Type: Integer

◆ kIsScheduled

constexpr const char* tvm::tirx::attr::kIsScheduled = "tirx.is_scheduled"
constexpr

Mark the function as scheduled, so the default schedule will pass will skip it.

Type: Integer

◆ kKernelLaunchParams

constexpr const char* tvm::tirx::attr::kKernelLaunchParams = "tirx.kernel_launch_params"
constexpr

List of thread IterVar that a DeviceLaunch function corresponds to.

Type: ffi::Array<ffi::String>

We call a device kernel launch function f using the following convention:

Call(f, [arg1, arg2, ..., arg_n, work_size_1, work_size_2, ... work_size_m, dyn_shmem_size])

Here n = len(arg), m = len(work_size) = len(launch_params)-1.

The list of kernel launch params indicates which additional parameters will be provided to the ffi::Function by the calling scope.

  • "threadIdx.x", "threadIdx.y", "threadIdx.z"

    The extent of the thread count in x/y/z, to be used when launching the compute kernel on the device. For example, the gridDimX/Y/Z parameters passed to cuLaunchKernel when launching a CUDA kernel, or the groupCountX/Y/Z parameters passed to vkCmdDispatch when dispatching a compute pipeline to Vulkan.

  • "blockIdx.x", "blockIdx.y", "blockIdx.z"

    The extent of the block iterators, to be used when launching the compute kernel on the device. For example, the blockDimX/Y/Z parameters passed to cuLaunchKernel when launching a CUDA kernel. For runtimes that do not require the block to be provided externally, this parameter is ignored. For example, the spv::ExecutionModeLocalSize for SPIR-V shaders on Vulkan, where this parameter is defined in the shader.

  • tvm::runtime::launch_param::kUseDynamicSharedMemoryTag

    The size of the shared memory that may be allocated internally by the kernel. For example, exposed as the CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES attribute in CUDA.

    Defined as "tirx.use_dyn_shared_memory".

See also
tvm::CallingConv::kDeviceKernelLaunch

◆ kNoAlias

constexpr const char* tvm::tirx::attr::kNoAlias = "tirx.noalias"
constexpr

Whether to set noalias rule on the function arguments.

Type: Integer

◆ kVolatile

constexpr const char* tvm::tirx::attr::kVolatile = "tirx.volatile"
constexpr

Annotation key on AllocBuffer marking the allocation as volatile.

◆ pragma_auto_unroll_max_step

constexpr const char* tvm::tirx::attr::pragma_auto_unroll_max_step = "pragma_auto_unroll_max_step"
constexpr

Pragma: auto-unroll, max_step.

◆ pragma_import_c

constexpr const char* tvm::tirx::attr::pragma_import_c = "pragma_import_c"
constexpr

Import C source or file into the final code gen module.

◆ pragma_import_llvm

constexpr const char* tvm::tirx::attr::pragma_import_llvm = "pragma_import_llvm"
constexpr

Import llvm source or file into the final code gen module.

◆ pragma_scope_prefix

constexpr const char* tvm::tirx::attr::pragma_scope_prefix = "pragma_"
constexpr

Mark region is guarded by the pragma extension.

◆ pragma_tensor_core

constexpr const char* tvm::tirx::attr::pragma_tensor_core = "pragma_tensor_core"
constexpr

Try to modify the AST to support Tensor Core.

◆ pragma_unroll_explicit

constexpr const char* tvm::tirx::attr::pragma_unroll_explicit = "pragma_unroll_explicit"
constexpr

Pragma: unroll explicit.

◆ storage_alignment

constexpr const char* tvm::tirx::attr::storage_alignment = "storage_alignment"
constexpr

Mark storage alignment requirement of buffers.

◆ thread_extent

constexpr const char* tvm::tirx::attr::thread_extent = "thread_extent"
constexpr

Mark launching extent of thread, used by device API.