25 #ifndef TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
26 #define TVM_RELAX_DISTRIBUTED_STRUCT_INFO_H_
32 namespace distributed {
49 refl::ObjectDef<PlacementSpecNode>()
54 static constexpr
const char*
_type_key =
"relax.distributed.PlacementSpec";
99 static constexpr
const char*
_type_key =
"relax.distributed.Placement";
109 TVM_DLL
explicit Placement(Array<PlacementSpec> dim_specs);
135 refl::ObjectDef<DTensorStructInfoNode>()
141 static constexpr
const char*
_type_key =
"relax.DTensorStructInfo";
Container of constant int that adds more constructors.
Definition: expr.h:612
Definition: source_map.h:113
Base type of all structure information.
Definition: expr.h:110
Managed reference to StructInfoNode.
Definition: expr.h:135
Managed reference to TensorStructInfoNode.
Definition: struct_info.h:196
StructInfo of DTensor (Distributed Tensor).
Definition: struct_info.h:118
TensorStructInfo tensor_sinfo
The struct info inherited from TensorStructInfo.
Definition: struct_info.h:123
Placement placement
The placement of the tensor among the device mesh.
Definition: struct_info.h:131
TVM_DECLARE_FINAL_OBJECT_INFO(DTensorStructInfoNode, StructInfoNode)
DeviceMesh device_mesh
The device mesh of the tensor.
Definition: struct_info.h:127
static void RegisterReflection()
Definition: struct_info.h:133
static constexpr const char * _type_key
Definition: struct_info.h:141
Managed reference to DTensorStructInfoNode.
Definition: struct_info.h:149
DTensorStructInfo(TensorStructInfo tensor_sinfo, DeviceMesh device_mesh, Placement placement, Span span=Span())
Construction with device mesh and placement.
TVM_DEFINE_OBJECT_REF_METHODS(DTensorStructInfo, StructInfo, DTensorStructInfoNode)
Managed reference to a DeviceMesh.
Definition: global_info.h:64
Describes how data is distributed in each dimension of the device mesh.
Definition: struct_info.h:86
TVM_DECLARE_FINAL_OBJECT_INFO(PlacementNode, Object)
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: struct_info.h:98
static constexpr const char * _type_key
Definition: struct_info.h:99
Array< PlacementSpec > dim_specs
specs for each dim of device mesh.
Definition: struct_info.h:89
static void RegisterReflection()
Definition: struct_info.h:93
Describes how data is distributed in one dimension of the device mesh.
Definition: struct_info.h:37
static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind
Definition: struct_info.h:55
PlacementSpecKind kind
The kind of placement spec. Possible values: kSharding and kReplica.
Definition: struct_info.h:45
TVM_DECLARE_BASE_OBJECT_INFO(PlacementSpecNode, Object)
static void RegisterReflection()
Definition: struct_info.h:47
static constexpr const char * _type_key
Definition: struct_info.h:54
int axis
If the kind is sharding, this value represents the tensor dimension to shard. otherwise,...
Definition: struct_info.h:42
Managed reference to PlacementSpecNode.
Definition: struct_info.h:63
TVM_DEFINE_OBJECT_REF_METHODS(PlacementSpec, ObjectRef, PlacementSpecNode)
static PlacementSpec Replica()
static PlacementSpec Sharding(int axis)
Managed reference to a Placement.
Definition: struct_info.h:107
static Placement FromText(String text_repr)
replica dim is printed as "R" and sharding dim is printed as "S[i]".]
Placement(Array< PlacementSpec > dim_specs)
TVM_DEFINE_OBJECT_REF_METHODS(Placement, ObjectRef, PlacementNode)
Definition: struct_info.h:72
TVM_DECLARE_FINAL_OBJECT_INFO(ShardingNode, PlacementSpecNode)
static void RegisterReflection()
Definition: struct_info.h:77
Integer sharding_dim
The dimension of tensor we shard.
Definition: struct_info.h:75
Definition: repr_printer.h:91
PlacementSpecKind
Definition: struct_info.h:34
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
Data structure for distributed inference.