device_api.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef TVM_RUNTIME_DEVICE_API_H_
#define TVM_RUNTIME_DEVICE_API_H_

#include <tvm/runtime/c_runtime_api.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/runtime/packed_func.h>

#include <string>

namespace tvm {
namespace runtime {
enum DeviceAttrKind : int {
  kExist = 0,
  kMaxThreadsPerBlock = 1,
  kWarpSize = 2,
  kMaxSharedMemoryPerBlock = 3,
  kComputeVersion = 4,
  kDeviceName = 5,
  kMaxClockRate = 6,
  kMultiProcessorCount = 7,
  kMaxThreadDimensions = 8,
  kMaxRegistersPerBlock = 9,
  kGcnArch = 10,
  kApiVersion = 11,
  kDriverVersion = 12
};

#ifdef TVM_KALLOC_ALIGNMENT

constexpr int kAllocAlignment = TVM_KALLOC_ALIGNMENT;

constexpr int kTempAllocaAlignment = TVM_KALLOC_ALIGNMENT;
#else

constexpr int kAllocAlignment = 64;

constexpr int kTempAllocaAlignment = 64;
#endif  // TVM_KALLOC_ALIGNMENT

constexpr int kMaxStackAlloca = 1024;

constexpr int kDefaultWorkspaceAlignment = 1;

class TVM_DLL DeviceAPI {
 public:
  virtual ~DeviceAPI() {}
  virtual void SetDevice(Device dev) = 0;
  virtual void GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) = 0;

  virtual void GetTargetProperty(Device dev, const std::string& property, TVMRetValue* rv) {}

  virtual void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment,
                               DLDataType type_hint) = 0;
  virtual void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType dtype,
                               Optional<String> mem_scope = NullOpt);
  virtual void FreeDataSpace(Device dev, void* ptr) = 0;
  virtual void CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle stream);
  virtual TVMStreamHandle CreateStream(Device dev);

  virtual void FreeStream(Device dev, TVMStreamHandle stream);

  virtual void StreamSync(Device dev, TVMStreamHandle stream) = 0;
  virtual void SetStream(Device dev, TVMStreamHandle stream) {}
  virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src, TVMStreamHandle event_dst);
  virtual void* AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint = {});
  virtual void FreeWorkspace(Device dev, void* ptr);

  static DeviceAPI* Get(Device dev, bool allow_missing = false);

  static bool NeedSetDevice(int device_type) {
    return device_type != kDLCPU && device_type != kDLMicroDev;
  }

 protected:
  virtual void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset,
                              size_t num_bytes, Device dev_from, Device dev_to,
                              DLDataType type_hint, TVMStreamHandle stream);
};

constexpr int kRPCSessMask = 128;
static_assert(kRPCSessMask >= TVMDeviceExtType_End);

inline const char* DeviceName(int type) {
  switch (type) {
    case kDLCPU:
      return "cpu";
    case kDLCUDA:
      return "cuda";
    case kDLCUDAHost:
      return "cuda_host";
    case kDLCUDAManaged:
      return "cuda_managed";
    case kDLOpenCL:
      return "opencl";
    case kDLSDAccel:
      return "sdaccel";
    case kDLAOCL:
      return "aocl";
    case kDLVulkan:
      return "vulkan";
    case kDLMetal:
      return "metal";
    case kDLVPI:
      return "vpi";
    case kDLROCM:
      return "rocm";
    case kDLROCMHost:
      return "rocm_host";
    case kDLExtDev:
      return "ext_dev";
    case kDLOneAPI:
      return "oneapi";
    case kDLWebGPU:
      return "webgpu";
    case kDLHexagon:
      return "hexagon";
    case kOpenGL:
      return "opengl";
    case kDLMicroDev:
      return "microdev";
    default:
      LOG(FATAL) << "unknown type =" << type;
  }
}

inline bool IsRPCSessionDevice(Device dev) { return (dev.device_type / kRPCSessMask) > 0; }

inline int GetRPCSessionIndex(Device dev) {
  ICHECK(IsRPCSessionDevice(dev)) << "GetRPCSessionIndex: dev has no RPC session";
  return dev.device_type / kRPCSessMask - 1;
}

inline Device RemoveRPCSessionMask(Device dev) {
  dev.device_type = static_cast<DLDeviceType>(dev.device_type % kRPCSessMask);
  return dev;
}

inline std::ostream& operator<<(std::ostream& os, DLDevice dev) {  // NOLINT(*)
  if (tvm::runtime::IsRPCSessionDevice(dev)) {
    os << "remote[" << tvm::runtime::GetRPCSessionIndex(dev) << "]-";
    dev = tvm::runtime::RemoveRPCSessionMask(dev);
  }
  os << tvm::runtime::DeviceName(static_cast<int>(dev.device_type)) << "(" << dev.device_id << ")";
  return os;
}

inline Device AddRPCSessionMask(Device dev, int session_table_index) {
  CHECK(!IsRPCSessionDevice(dev)) << "AddRPCSessionMask: dev already non-zero RPCSessionIndex: "
                                  << dev;
  dev.device_type =
      static_cast<DLDeviceType>(dev.device_type | (kRPCSessMask * (session_table_index + 1)));
  return dev;
}

}  // namespace runtime
}  // namespace tvm

#endif  // TVM_RUNTIME_DEVICE_API_H_