experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allow tracking cuLaunchCooperativeKernel cbid.

Browse Source 
CUPTI_ACTIVITY_KIND_KERNEL will include this in activity API, if not tracing the launch, we are missing annotation for these kind of kernels.
For the event pairs solution, such kernels will be ignored if not tracked; causing confusing between two methods.

PiperOrigin-RevId: 278706189
Change-Id: I29487ebc737beca439ef9e9332c2dcf2d3e5c5b7
This commit is contained in:
A. Unique TensorFlower 2019-11-05 14:15:01 -08:00 committed by TensorFlower Gardener
parent 60b7d06456
commit 97ff8b57a7
1 changed files with 17 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tensorflow/core/profiler/internal/gpu/cupti_tracer.cc
View File

@ -676,9 +676,9 @@ class CudaEventRecorder {
// Registers the start of a kernel launch. The returned index should be passed
// to StopKernel() after the kernel launch has completed.
template <typename T>
size_t StartKernel(const char *kernel_name, CUcontext context,
uint32 correlation_id,
const cuLaunchKernel_params *params) {
uint32 correlation_id, const T *params) {
CUstream stream = params->hStream;
KernelRecord record = {kernel_name, context, stream, correlation_id};
record.details.registers_per_thread = 0; // unknown.
@ -968,10 +968,20 @@ class CuptiDriverApiHookWithCudaEvent : public CuptiDriverApiHook {
DCHECK_NE(cbdata->symbolName, nullptr);
auto params =
static_cast<const cuLaunchKernel_params *>(cbdata->functionParams);
*cbdata->correlationData = recorder->StartKernel(
*cbdata->correlationData = recorder->StartKernel<cuLaunchKernel_params>(
cbdata->symbolName, cbdata->context, cbdata->correlationId, params);
break;
}
case CUPTI_DRIVER_TRACE_CBID_cuLaunchCooperativeKernel: {
DCHECK_NE(cbdata->symbolName, nullptr);
auto params = static_cast<const cuLaunchCooperativeKernel_params_st *>(
cbdata->functionParams);
*cbdata->correlationData =
recorder->StartKernel<cuLaunchCooperativeKernel_params_st>(
cbdata->symbolName, cbdata->context, cbdata->correlationId,
params);
break;
}
case CUPTI_DRIVER_TRACE_CBID_cuMemcpy: {
auto params =
static_cast<const cuMemcpy_params *>(cbdata->functionParams);
@ -1010,6 +1020,10 @@ class CuptiDriverApiHookWithCudaEvent : public CuptiDriverApiHook {
StartMemcpyAsync<cuMemcpyDtoDAsync_v2_params>(
CuptiTracerEventType::MemcpyD2D, cbdata, recorder);
break;
case CUPTI_DRIVER_TRACE_CBID_cuLaunchCooperativeKernelMultiDevice:
// TODO: track these kind of events.
VLOG(1) << "untracked cuLaunchCooperativeKernelMultiDevice";
break;
default:
VLOG(1) << "Unexpected callback id: " << cbid;
break;