Skip to content

Commit

Permalink
proper device query through rocm api
Browse files Browse the repository at this point in the history
  • Loading branch information
@petrex
petrex committed Nov 15, 2019
1 parent 3f6b3db commit 60064ec
Showing 1 changed file with 123 additions and 102 deletions.
225 changes: 123 additions & 102 deletions src/runtime/rocm/rocm_device_api.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@
* 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
Expand All @@ -22,14 +22,12 @@
* \file rocm_device_api.cc
* \brief GPU specific API
*/
#include <tvm/runtime/device_api.h>

#include <dmlc/logging.h>
#include <dmlc/thread_local.h>
#include <hip/hip_runtime_api.h>
#include <hsa/hsa.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/registry.h>
#include "../../../include/tvm/runtime/device_api.h"
#include "rocm_common.h"

namespace tvm {
Expand All @@ -55,130 +53,153 @@ class ROCMDeviceAPI final : public DeviceAPI {
break;
}
case kMaxThreadsPerBlock: {
value = 1024;
ROCM_CALL(hipDeviceGetAttribute(
&value, hipDeviceAttributeMaxThreadsPerBlock, ctx.device_id));
break;
}
case kWarpSize: {
value = 64;
ROCM_CALL(hipDeviceGetAttribute(&value, hipDeviceAttributeWarpSize,
ctx.device_id));
break;
}
case kMaxSharedMemoryPerBlock: {
ROCM_CALL(hipDeviceGetAttribute(
&value, hipDeviceAttributeMaxSharedMemoryPerBlock, ctx.device_id));
break;
}
case kComputeVersion: {
std::ostringstream os;
ROCM_CALL(hipDeviceGetAttribute(
&value, hipDeviceAttributeComputeCapabilityMajor, ctx.device_id));
os << value << ".";
ROCM_CALL(hipDeviceGetAttribute(
&value, hipDeviceAttributeComputeCapabilityMinor, ctx.device_id));
os << value;
*rv = os.str();
return;
}
case kDeviceName:
return;
case kMaxClockRate: {
ROCM_CALL(hipDeviceGetAttribute(&value, hipDeviceAttributeClockRate,
ctx.device_id));
break;
}
case kMaxSharedMemoryPerBlock: return;
case kComputeVersion:
case kDeviceName: return;
case kMaxClockRate: return;
case kMultiProcessorCount: return;
case kMaxThreadDimensions: return;
case kMultiProcessorCount: {
ROCM_CALL(hipDeviceGetAttribute(
&value, hipDeviceAttributeMultiprocessorCount, ctx.device_id));
break;
}
case kMaxThreadDimensions: {
int dims[3];
ROCM_CALL(hipDeviceGetAttribute(
&dims[0], hipDeviceAttributeMaxBlockDimX, ctx.device_id));
ROCM_CALL(hipDeviceGetAttribute(
&dims[1], hipDeviceAttributeMaxBlockDimY, ctx.device_id));
ROCM_CALL(hipDeviceGetAttribute(
&dims[2], hipDeviceAttributeMaxBlockDimZ, ctx.device_id));

std::stringstream ss;
ss << "[" << dims[0] << ", " << dims[1] << ", " << dims[2] << "]";
*rv = ss.str();
return;
}
case kGcnArch: {
hipDeviceProp_t prop;
ROCM_CALL(hipGetDeviceProperties(&prop, ctx.device_id));
*rv = prop.gcnArch;
return;
}
*rv = value;
}
void* AllocDataSpace(TVMContext ctx, size_t nbytes, size_t alignment,
TVMType type_hint) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
CHECK_EQ(256 % alignment, 0U) << "ROCM space is aligned at 256 bytes";
void* ret;
ROCM_CALL(hipMalloc(&ret, nbytes));
return ret;
}
*rv = value;
}
void* AllocDataSpace(TVMContext ctx,
size_t nbytes,
size_t alignment,
TVMType type_hint) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
CHECK_EQ(256 % alignment, 0U)
<< "ROCM space is aligned at 256 bytes";
void *ret;
ROCM_CALL(hipMalloc(&ret, nbytes));
return ret;
}

void FreeDataSpace(TVMContext ctx, void* ptr) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
ROCM_CALL(hipFree(ptr));
}
void FreeDataSpace(TVMContext ctx, void* ptr) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
ROCM_CALL(hipFree(ptr));
}

void CopyDataFromTo(const void* from,
size_t from_offset,
void* to,
size_t to_offset,
size_t size,
TVMContext ctx_from,
TVMContext ctx_to,
TVMType type_hint,
TVMStreamHandle stream) final {
hipStream_t hip_stream = static_cast<hipStream_t>(stream);
from = static_cast<const char*>(from) + from_offset;
to = static_cast<char*>(to) + to_offset;
if (ctx_from.device_type == kDLROCM && ctx_to.device_type == kDLROCM) {
ROCM_CALL(hipSetDevice(ctx_from.device_id));
if (ctx_from.device_id == ctx_to.device_id) {
GPUCopy(from, to, size, hipMemcpyDeviceToDevice, hip_stream);
void CopyDataFromTo(const void* from, size_t from_offset, void* to,
size_t to_offset, size_t size, TVMContext ctx_from,
TVMContext ctx_to, TVMType type_hint,
TVMStreamHandle stream) final {
hipStream_t hip_stream = static_cast<hipStream_t>(stream);
from = static_cast<const char*>(from) + from_offset;
to = static_cast<char*>(to) + to_offset;
if (ctx_from.device_type == kDLROCM && ctx_to.device_type == kDLROCM) {
ROCM_CALL(hipSetDevice(ctx_from.device_id));
if (ctx_from.device_id == ctx_to.device_id) {
GPUCopy(from, to, size, hipMemcpyDeviceToDevice, hip_stream);
} else {
hipMemcpyPeerAsync(to, ctx_to.device_id, from, ctx_from.device_id,
size, hip_stream);
}
} else if (ctx_from.device_type == kDLROCM &&
ctx_to.device_type == kDLCPU) {
ROCM_CALL(hipSetDevice(ctx_from.device_id));
GPUCopy(from, to, size, hipMemcpyDeviceToHost, hip_stream);
} else if (ctx_from.device_type == kDLCPU &&
ctx_to.device_type == kDLROCM) {
ROCM_CALL(hipSetDevice(ctx_to.device_id));
GPUCopy(from, to, size, hipMemcpyHostToDevice, hip_stream);
} else {
hipMemcpyPeerAsync(to, ctx_to.device_id,
from, ctx_from.device_id,
size, hip_stream);
LOG(FATAL) << "expect copy from/to GPU or between GPU";
}
} else if (ctx_from.device_type == kDLROCM && ctx_to.device_type == kDLCPU) {
ROCM_CALL(hipSetDevice(ctx_from.device_id));
GPUCopy(from, to, size, hipMemcpyDeviceToHost, hip_stream);
} else if (ctx_from.device_type == kDLCPU && ctx_to.device_type == kDLROCM) {
ROCM_CALL(hipSetDevice(ctx_to.device_id));
GPUCopy(from, to, size, hipMemcpyHostToDevice, hip_stream);
} else {
LOG(FATAL) << "expect copy from/to GPU or between GPU";
}
}

void StreamSync(TVMContext ctx, TVMStreamHandle stream) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
ROCM_CALL(hipStreamSynchronize(static_cast<hipStream_t>(stream)));
}

void SetStream(TVMContext ctx, TVMStreamHandle stream) final {
ROCMThreadEntry::ThreadLocal()
->stream = static_cast<hipStream_t>(stream);
}
void StreamSync(TVMContext ctx, TVMStreamHandle stream) final {
ROCM_CALL(hipSetDevice(ctx.device_id));
ROCM_CALL(hipStreamSynchronize(static_cast<hipStream_t>(stream)));
}

void* AllocWorkspace(TVMContext ctx, size_t size, TVMType type_hint) final {
return ROCMThreadEntry::ThreadLocal()->pool.AllocWorkspace(ctx, size);
}
void SetStream(TVMContext ctx, TVMStreamHandle stream) final {
ROCMThreadEntry::ThreadLocal()->stream = static_cast<hipStream_t>(stream);
}

void FreeWorkspace(TVMContext ctx, void* data) final {
ROCMThreadEntry::ThreadLocal()->pool.FreeWorkspace(ctx, data);
}
void* AllocWorkspace(TVMContext ctx, size_t size, TVMType type_hint) final {
return ROCMThreadEntry::ThreadLocal()->pool.AllocWorkspace(ctx, size);
}

static const std::shared_ptr<ROCMDeviceAPI>& Global() {
static std::shared_ptr<ROCMDeviceAPI> inst =
std::make_shared<ROCMDeviceAPI>();
return inst;
}
void FreeWorkspace(TVMContext ctx, void* data) final {
ROCMThreadEntry::ThreadLocal()->pool.FreeWorkspace(ctx, data);
}

private:
static void GPUCopy(const void* from,
void* to,
size_t size,
hipMemcpyKind kind,
hipStream_t stream) {
if (stream != 0) {
ROCM_CALL(hipMemcpyAsync(to, from, size, kind, stream));
} else {
ROCM_CALL(hipMemcpy(to, from, size, kind));
static const std::shared_ptr<ROCMDeviceAPI>& Global() {
static std::shared_ptr<ROCMDeviceAPI> inst =
std::make_shared<ROCMDeviceAPI>();
return inst;
}
}
};

typedef dmlc::ThreadLocalStore<ROCMThreadEntry> ROCMThreadStore;
private:
static void GPUCopy(const void* from, void* to, size_t size,
hipMemcpyKind kind, hipStream_t stream) {
if (stream != 0) {
ROCM_CALL(hipMemcpyAsync(to, from, size, kind, stream));
} else {
ROCM_CALL(hipMemcpy(to, from, size, kind));
}
}
};

ROCMThreadEntry::ROCMThreadEntry()
: pool(kDLROCM, ROCMDeviceAPI::Global()) {
}
typedef dmlc::ThreadLocalStore<ROCMThreadEntry> ROCMThreadStore;

ROCMThreadEntry* ROCMThreadEntry::ThreadLocal() {
return ROCMThreadStore::Get();
}
ROCMThreadEntry::ROCMThreadEntry() : pool(kDLROCM, ROCMDeviceAPI::Global()) {}

TVM_REGISTER_GLOBAL("device_api.rocm")
.set_body([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = ROCMDeviceAPI::Global().get();
*rv = static_cast<void*>(ptr);
});
ROCMThreadEntry* ROCMThreadEntry::ThreadLocal() {
return ROCMThreadStore::Get();
}

TVM_REGISTER_GLOBAL("device_api.rocm")
.set_body([](TVMArgs args, TVMRetValue* rv) {
DeviceAPI* ptr = ROCMDeviceAPI::Global().get();
*rv = static_cast<void*>(ptr);
});
} // namespace runtime
} // namespace tvm

0 comments on commit 60064ec

Please sign in to comment.