From e8a37135a338ec03126f49d8da554544d1a164be Mon Sep 17 00:00:00 2001
From: Aaron Shi <aaronshi@meta.com>
Date: Tue, 5 Dec 2023 06:31:39 -0800
Subject: [PATCH] Specify cuContext in range profiler structs (#643)

Summary:
Pull Request resolved: https://github.com/pytorch/kineto/pull/643

## Context:
I was experiencing some crashes with the range proiler. The backtrace showed ::disable() failing and I found that the `.ctx` member of the configuration struct was not set.

## Changes:
This diff specifies cuContext for all structs that have it as a required member.

This diff also specifies the `.targetNestingLevel` member for the `setConfigParams` struct as it was unset.

Reviewed By: benghaem

Differential Revision: D38435331

fbshipit-source-id: 3a2a297f7ed5665ed3d58df7f48ae8a64e5ac956
---
 libkineto/src/CuptiRangeProfilerApi.cpp | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/libkineto/src/CuptiRangeProfilerApi.cpp b/libkineto/src/CuptiRangeProfilerApi.cpp
index 4b3963b08..25bc32dd3 100644
--- a/libkineto/src/CuptiRangeProfilerApi.cpp
+++ b/libkineto/src/CuptiRangeProfilerApi.cpp
@@ -388,7 +388,9 @@ CuptiRBProfilerSession::CuptiRBProfilerSession(
             << counterDataScratchBuffer.size()  << " B";
 
   beginPassParams_ = {CUpti_Profiler_BeginPass_Params_STRUCT_SIZE, nullptr};
+  beginPassParams_.ctx = cuContext_;
   endPassParams_ = {CUpti_Profiler_EndPass_Params_STRUCT_SIZE, nullptr};
+  endPassParams_.ctx = cuContext_;
 
   initSuccess_ = true;
   profiler_map[deviceId_] = this;
@@ -448,9 +450,10 @@ void CuptiRBProfilerSession::startInternal(
   setConfigParams.ctx = cuContext_;
   setConfigParams.pConfig = configImage.data();
   setConfigParams.configSize = configImage.size();
-  setConfigParams.passIndex = 0;
   setConfigParams.minNestingLevel = 1;
   setConfigParams.numNestingLevels = numNestingLevels_;
+  setConfigParams.passIndex = 0;
+  setConfigParams.targetNestingLevel = setConfigParams.minNestingLevel;
   status = CUPTI_CALL(cuptiProfilerSetConfig(&setConfigParams));
 
   if (status != CUPTI_SUCCESS) {
@@ -475,10 +478,12 @@ void CuptiRBProfilerSession::stop() {
 
   CUpti_Profiler_UnsetConfig_Params unsetConfigParams = {
       CUpti_Profiler_UnsetConfig_Params_STRUCT_SIZE, nullptr};
+  unsetConfigParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerUnsetConfig(&unsetConfigParams));
 
   CUpti_Profiler_EndSession_Params endSessionParams = {
       CUpti_Profiler_EndSession_Params_STRUCT_SIZE, nullptr};
+  endSessionParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerEndSession(&endSessionParams));
 
   disableKernelCallbacks();
@@ -508,6 +513,7 @@ void CuptiRBProfilerSession::flushCounterData() {
   LOG(INFO) << "Flushing counter data on device = " << deviceId_;
   CUpti_Profiler_FlushCounterData_Params flushCounterDataParams = {
       CUpti_Profiler_FlushCounterData_Params_STRUCT_SIZE, nullptr};
+  flushCounterDataParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerFlushCounterData(&flushCounterDataParams));
 }
 
@@ -519,6 +525,7 @@ void CuptiRBProfilerSession::enable() {
   }
   CUpti_Profiler_EnableProfiling_Params enableProfilingParams = {
       CUpti_Profiler_EnableProfiling_Params_STRUCT_SIZE, nullptr};
+  enableProfilingParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerEnableProfiling(&enableProfilingParams));
 }
 
@@ -529,6 +536,7 @@ void CuptiRBProfilerSession::disable() {
   }
   CUpti_Profiler_DisableProfiling_Params disableProfilingParams = {
       CUpti_Profiler_DisableProfiling_Params_STRUCT_SIZE, nullptr};
+  disableProfilingParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerDisableProfiling(&disableProfilingParams));
 }
 
@@ -537,6 +545,7 @@ void CuptiRBProfilerSession::pushRange(const std::string& rangeName) {
   LOG(INFO) << " CUPTI pushrange ( " << rangeName << " )";
   CUpti_Profiler_PushRange_Params pushRangeParams = {
       CUpti_Profiler_PushRange_Params_STRUCT_SIZE, nullptr};
+  pushRangeParams.ctx = cuContext_;
   pushRangeParams.pRangeName = rangeName.c_str();
   CUPTI_CALL(cuptiProfilerPushRange(&pushRangeParams));
 }
@@ -545,6 +554,7 @@ void CuptiRBProfilerSession::popRange() {
   LOG(INFO) << " CUPTI pop range";
   CUpti_Profiler_PopRange_Params popRangeParams = {
       CUpti_Profiler_PopRange_Params_STRUCT_SIZE, nullptr};
+  popRangeParams.ctx = cuContext_;
   CUPTI_CALL(cuptiProfilerPopRange(&popRangeParams));
 }