tvm
|
Attributes for the "on_device" special operator. More...
#include <annotation.h>
Public Member Functions | |
TVM_DECLARE_ATTRS (OnDeviceAttrs, "relay.attrs.OnDeviceAttrs") | |
![]() | |
void | VisitAttrs (AttrVisitor *v) |
void | VisitNonDefaultAttrs (AttrVisitor *v) |
Visit attributes that do not equal the default value. More... | |
void | InitByPackedArgs (const runtime::TVMArgs &args, bool allow_unknown) final |
Initialize the attributes by arguments. More... | |
bool | SEqualReduce (const OnDeviceAttrs *other, SEqualReducer equal) const |
void | SHashReduce (SHashReducer hash_reducer) const |
Array< AttrFieldInfo > | ListFieldInfo () const final |
Get the field information. More... | |
![]() | |
virtual | ~BaseAttrsNode () |
virtual destructor More... | |
template<typename... Args> | |
void | InitBySeq (Args &&... args) |
Initialize the attributes by sequence of arguments. More... | |
void | PrintDocString (std::ostream &os) const |
Print readible docstring to ostream, add newline. More... | |
TVM_DECLARE_BASE_OBJECT_INFO (BaseAttrsNode, Object) | |
![]() | |
uint32_t | type_index () const |
std::string | GetTypeKey () const |
size_t | GetTypeKeyHash () const |
template<typename TargetType > | |
bool | IsInstance () const |
bool | unique () const |
Object () | |
Object (const Object &other) | |
Object (Object &&other) | |
Object & | operator= (const Object &other) |
Object & | operator= (Object &&other) |
Public Attributes | |
int | device_type = kInvalidDeviceType |
Device type on which argument expression should be evaluated. More... | |
bool | is_fixed = false |
If true, the result device must also be device_type and device planning should not insert any "device_copy" calls to respect this annotation. More... | |
Additional Inherited Members | |
![]() | |
using | TVMArgs = runtime::TVMArgs |
using | TVMRetValue = runtime::TVMRetValue |
![]() | |
typedef void(* | FDeleter) (Object *self) |
Object deleter. More... | |
using | RefCounterType = std::atomic< int32_t > |
![]() | |
static std::string | TypeIndex2Key (uint32_t tindex) |
Get the type key of the corresponding index from runtime. More... | |
static size_t | TypeIndex2KeyHash (uint32_t tindex) |
Get the type key hash of the corresponding index from runtime. More... | |
static uint32_t | TypeKey2Index (const std::string &key) |
Get the type index of the corresponding key from runtime. More... | |
static uint32_t | _GetOrAllocRuntimeTypeIndex () |
static uint32_t | RuntimeTypeIndex () |
![]() | |
static constexpr const bool | _type_has_method_sequal_reduce = true |
static constexpr const bool | _type_has_method_shash_reduce = true |
static constexpr const char * | _type_key = "Attrs" |
![]() | |
static constexpr const char * | _type_key = "runtime.Object" |
static constexpr bool | _type_final = false |
static constexpr uint32_t | _type_child_slots = 0 |
static constexpr bool | _type_child_slots_can_overflow = true |
static constexpr bool | _type_has_method_visit_attrs = true |
static constexpr bool | _type_has_method_sequal_reduce = false |
static constexpr bool | _type_has_method_shash_reduce = false |
static constexpr uint32_t | _type_index = TypeIndex::kDynamic |
![]() | |
void | IncRef () |
developer function, increases reference counter. More... | |
void | DecRef () |
developer function, decrease reference counter. More... | |
![]() | |
static uint32_t | GetOrAllocRuntimeTypeIndex (const std::string &key, uint32_t static_tindex, uint32_t parent_tindex, uint32_t type_child_slots, bool type_child_slots_can_overflow) |
Get the type index using type key. More... | |
![]() | |
uint32_t | type_index_ {0} |
Type index(tag) that indicates the type of the object. More... | |
RefCounterType | ref_counter_ {0} |
The internal reference counter. More... | |
FDeleter | deleter_ = nullptr |
deleter of this object to enable customized allocation. If the deleter is nullptr, no deletion will be performed. The creator of the object must always set the deleter field properly. More... | |
Attributes for the "on_device" special operator.
The Relay call (aka 'annotation'):
constrains sub_expr
to execute and store its result on a device with DLDeviceType
2
(i.e. a kDLCuda
device). However the annotation itself may appear in an expression to be executed and stored on a different device. If so the compiler will automatically insert a "device_copy" call to mediate the transition between devices.
E.g.: Assuming x and y reside on the GPU and z on the CPU then:
indicates the add
should execute on the GPU but the multiply
should execute on the CPU. The compiler will rewrite this to:
The Relay call
is similar to the above, however the annotation itself must appear in an expression on the same device. The compiler will check the devices are consistent, and will not insert any "device_copy" call. This form of annotation shouldn't be necessary in user programs. However it is needed by the PlanDevices
pass to fully specify the results of device planning so that the pass is idempotent.
E.g.: The following program is equivalent to the above:
The "on_device" annotation with is_fixed=True
indicates unambiguously that a
is stored on the GPU.
|
inline |
int tvm::relay::OnDeviceAttrs::device_type = kInvalidDeviceType |
Device type on which argument expression should be evaluated.
bool tvm::relay::OnDeviceAttrs::is_fixed = false |
If true, the result device must also be device_type
and device planning should not insert any "device_copy" calls to respect this annotation.
This is used by the device planning pass itself when annotating the planned program.