tvm
virtual_device.h
Go to the documentation of this file.
1 /*
2  * Licensed to the Apache Software Foundation (ASF) under one
3  * or more contributor license agreements. See the NOTICE file
4  * distributed with this work for additional information
5  * regarding copyright ownership. The ASF licenses this file
6  * to you under the Apache License, Version 2.0 (the
7  * "License"); you may not use this file except in compliance
8  * with the License. You may obtain a copy of the License at
9  *
10  * http://www.apache.org/licenses/LICENSE-2.0
11  *
12  * Unless required by applicable law or agreed to in writing,
13  * software distributed under the License is distributed on an
14  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15  * KIND, either express or implied. See the License for the
16  * specific language governing permissions and limitations
17  * under the License.
18  */
19 
26 #ifndef TVM_TARGET_VIRTUAL_DEVICE_H_
27 #define TVM_TARGET_VIRTUAL_DEVICE_H_
28 
29 #include <tvm/ir/transform.h>
30 #include <tvm/target/target.h>
31 
32 #include <string>
33 #include <unordered_set>
34 #include <utility>
35 
36 namespace tvm {
37 
45 using MemoryScope = ffi::String;
46 
47 // NOTE: cannot use enum as they are out of bound of the original enum
48 // and results in an undefined behavior
49 // A 'null' device type, does not correspond to any DLDeviceType enum.
50 // TODO(mbs): This is to help us as we transition away from representing the 'homogenous' case
51 // as a singleton target map indexed by the invalid DLDeviceType '0'.
52 constexpr int kNullDeviceType = 0;
53 
54 // An 'invalid' device type, does not correspond to any DLDeviceType enum.
55 constexpr int kInvalidDeviceType = -1;
56 
172 class VirtualDeviceNode : public AttrsNodeReflAdapter<VirtualDeviceNode> {
173  private:
185  int /* actually DLDeviceType */ device_type_int;
186 
187  public:
188  DLDeviceType device_type() const { return static_cast<DLDeviceType>(device_type_int); }
189 
197 
206 
213 
218  bool IsFullyUnconstrained() const {
219  return !target.defined() && device_type() == kInvalidDeviceType && virtual_device_id == -1 &&
220  memory_scope.empty();
221  }
222 
227  bool IsFullyConstrained() const {
228  return target.defined() && virtual_device_id != -1 && !memory_scope.empty();
229  }
230 
237  Device ToDevice() const {
238  ICHECK(device_type_int != kInvalidDeviceType);
239  ICHECK(virtual_device_id != -1);
240  Device device;
241  device.device_type = device_type();
242  device.device_id = virtual_device_id;
243  return device;
244  }
245 
246  static void RegisterReflection() {
247  namespace refl = tvm::ffi::reflection;
248  refl::ObjectDef<VirtualDeviceNode>()
249  .def_ro("device_type_int", &VirtualDeviceNode::device_type_int,
250  "The type of the virtual device.", refl::DefaultValue(kInvalidDeviceType))
251  .def_ro("virtual_device_id", &VirtualDeviceNode::virtual_device_id,
252  "The device id of the virtual device.", refl::DefaultValue(-1))
253  .def_ro("target", &VirtualDeviceNode::target,
254  "The target describing how to compile for the virtual device.",
255  refl::DefaultValue(Target()))
256  .def_ro("memory_scope", &VirtualDeviceNode::memory_scope,
257  "The area of memory w.r.t. the virtual device where data is stored.",
258  refl::DefaultValue(""));
259  }
261 
262  friend class VirtualDevice;
263 };
264 
268 class VirtualDevice : public ObjectRef {
269  public:
280  explicit VirtualDevice(int device_type_int = kInvalidDeviceType, int virtual_device_id = -1,
281  Target target = {}, MemoryScope memory_scope = {});
282 
285 
290  static VirtualDevice ForDeviceType(DLDeviceType device_type, int virtual_device_id = -1) {
291  ICHECK_GT(device_type, 0);
292  return VirtualDevice(device_type, virtual_device_id);
293  }
294  static VirtualDevice ForDeviceType(int device_type, int virtual_device_id = -1) {
295  return ForDeviceType(static_cast<DLDeviceType>(device_type), virtual_device_id);
296  }
297  static VirtualDevice ForDeviceType(const Integer& device_type, int virtual_device_id = -1) {
298  return ForDeviceType(static_cast<int>(device_type->value), virtual_device_id);
299  }
300 
302  static VirtualDevice ForDevice(const Device& device) {
303  return ForDeviceType(device.device_type, device.device_id);
304  }
305 
307  static VirtualDevice ForDeviceAndTarget(const Device& device, Target target) {
308  return VirtualDevice(device.device_type, device.device_id, std::move(target));
309  }
310 
312  static VirtualDevice ForTarget(Target target) {
313  DLDeviceType device_type = static_cast<DLDeviceType>(target->GetTargetDeviceType());
314  return VirtualDevice(device_type, /*virtual_device_id=*/0, std::move(target));
315  }
316 
318  static VirtualDevice ForMemoryScope(MemoryScope memory_scope) {
319  return VirtualDevice(kInvalidDeviceType, -1, {}, std::move(memory_scope));
320  }
321 
323  TVM_DLL static VirtualDevice ForDeviceTargetAndMemoryScope(const Device& device, Target target,
324  MemoryScope memory_scope) {
325  return VirtualDevice(device.device_type, device.device_id, std::move(target),
326  std::move(memory_scope));
327  }
328 
334  static ffi::Optional<VirtualDevice> Join(const VirtualDevice& lhs, const VirtualDevice& rhs);
335 
340  static VirtualDevice Default(const VirtualDevice& lhs, const VirtualDevice& rhs);
341 
343 
344  friend class VirtualDeviceCache; // Private implementation helper.
345 };
346 
356  public:
358  VirtualDevice Make(int device_type = kInvalidDeviceType, int virtual_device_id = -1,
359  Target target = {}, MemoryScope memory_scope = {});
360 
364  VirtualDevice Unique(const VirtualDevice& virtual_device);
365 
366  private:
368  std::unordered_set<VirtualDevice, StructuralHash, StructuralEqual> cache_;
369 };
370 
376 constexpr const char* kVirtualDevice = "virtual_device";
377 
378 } // namespace tvm
379 
380 #endif // TVM_TARGET_VIRTUAL_DEVICE_H_
Adapter for AttrsNode with the new reflection API.
Definition: attrs.h:385
Base class of all attribute class.
Definition: attrs.h:102
Container of constant int that adds more constructors.
Definition: expr.h:600
Managed reference class to TargetNode.
Definition: target.h:192
A cache of VirtualDevices. This can be used:
Definition: virtual_device.h:355
VirtualDevice Unique(const VirtualDevice &virtual_device)
Returns the unique VirtualDevice structurally equal to the given virtual_device.
VirtualDevice Make(int device_type=kInvalidDeviceType, int virtual_device_id=-1, Target target={}, MemoryScope memory_scope={})
Returns the unique VirtualDevice representing given fields.
Describes at compile time the constraints on where data is to be stored at runtime down to the (virtu...
Definition: virtual_device.h:172
static void RegisterReflection()
Definition: virtual_device.h:246
DLDeviceType device_type() const
Definition: virtual_device.h:188
Target target
The Target describing how to compile for the virtual device.
Definition: virtual_device.h:205
int virtual_device_id
The device identifier for the virtual device. This must be resolved to a physical device identifier e...
Definition: virtual_device.h:196
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("target.VirtualDevice", VirtualDeviceNode, BaseAttrsNode)
MemoryScope memory_scope
The scope of memory w.r.t. the virtual device which holds data.
Definition: virtual_device.h:212
Device ToDevice() const
Returns the (virtual) Device implied by this VirtualDevice. Both the device_type and virtual_device_m...
Definition: virtual_device.h:237
bool IsFullyUnconstrained() const
Returns true if virtual device is 'fully unconstrained', ie no target/device type,...
Definition: virtual_device.h:218
bool IsFullyConstrained() const
Returns true if virtual device is 'fully constrained', ie target, device id and memory scope are all ...
Definition: virtual_device.h:227
Managed reference class to VirtualDeviceNode.
Definition: virtual_device.h:268
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(VirtualDevice, ObjectRef, VirtualDeviceNode)
VirtualDevice(int device_type_int=kInvalidDeviceType, int virtual_device_id=-1, Target target={}, MemoryScope memory_scope={})
Construct a virtual device.
static VirtualDevice ForDeviceTargetAndMemoryScope(const Device &device, Target target, MemoryScope memory_scope)
Returns the VirtualDevice for device, target and memory_scope.
Definition: virtual_device.h:323
static ffi::Optional< VirtualDevice > Join(const VirtualDevice &lhs, const VirtualDevice &rhs)
Returns the 'join' of lhs and rhs. The result will agree pointwise with lhs and rhs on all their cons...
static VirtualDevice ForMemoryScope(MemoryScope memory_scope)
Returns the VirtualDevice for memory_scope alone.
Definition: virtual_device.h:318
static VirtualDevice Default(const VirtualDevice &lhs, const VirtualDevice &rhs)
Returns the 'default' of lhs and rhs. The result will be lhs, except any unconstrained fields in lhs ...
static VirtualDevice FullyUnconstrained()
Returns the unique fully unconstrained VirtualDevice.
static VirtualDevice ForDeviceAndTarget(const Device &device, Target target)
Returns the VirtualDevice for device and target.
Definition: virtual_device.h:307
static VirtualDevice ForDevice(const Device &device)
Returns the VirtualDevice for device.
Definition: virtual_device.h:302
static VirtualDevice ForDeviceType(const Integer &device_type, int virtual_device_id=-1)
Definition: virtual_device.h:297
static VirtualDevice ForDeviceType(DLDeviceType device_type, int virtual_device_id=-1)
Returns the VirtualDevice for device_type and (if not -1) virtual_device_id. The target and memory sc...
Definition: virtual_device.h:290
static VirtualDevice ForDeviceType(int device_type, int virtual_device_id=-1)
Definition: virtual_device.h:294
static VirtualDevice ForTarget(Target target)
Returns the VirtualDevice for target.
Definition: virtual_device.h:312
Definition: repr_printer.h:91
constexpr const char * device_type
The device type.
Definition: stmt.h:1063
Performance counters for profiling via the PAPI library.
Definition: analyzer.h:37
constexpr const char * kVirtualDevice
Definition: virtual_device.h:376
constexpr int kInvalidDeviceType
Definition: virtual_device.h:55
ffi::String MemoryScope
Abstract label for an area of memory.
Definition: global_info.h:37
DLDevice Device
Definition: device_api.h:42
constexpr int kNullDeviceType
Definition: virtual_device.h:52
Compilation target object.