[Memory Snapshot][BE] Clean up record function callback scope (pytorc…

…h#130265)

Summary: We can directly set the scope to at::RecordScope::USER_SCOPE for the at::RecordFunctionCallback object, rather than performing a check inside of the callback.

Test Plan:
Ran locally, works fine.

https://www.internalfb.com/pytorch_memory_visualizer/mvai_gpu_traces/tree/gpu_snapshot/fire-aaronshi-20240704-1709-7a80b83b/0/rank-0_itrn-1503.Jul_04_17_24_02.3577.snapshot.pickle

Differential Revision: D59477046

Pulled By: aaronenyeshi

Pull Request resolved: pytorch#130265
Approved by: https://github.com/davidberard98

torch/csrc/cuda/memory_snapshot.cpp

@@ -101,27 +101,23 @@ void _initRecordAnnotations() {
   static c10::once_flag ra_init;
   c10::call_once(ra_init, [&] {
     // Save user annotations to CCA memory snapshot tool
-    at::addThreadLocalCallback(at::RecordFunctionCallback(
-        [](const at::RecordFunction& fn)
-            -> std::unique_ptr<at::ObserverContext> {
-          if (fn.scope() != at::RecordScope::USER_SCOPE) {
-            return nullptr; // only record user-defined scopes.
-          }
-          unwind::Frame frame{fn.name(), "START", 0};
-          auto r = std::make_shared<CapturedTraceback>();
-          r->recordUserDefinedFrame(frame);
-          c10::cuda::CUDACachingAllocator::recordAnnotation(r);
-          return nullptr;
-        },
-        [](const at::RecordFunction& fn, at::ObserverContext* ctx_ptr) {
-          if (fn.scope() != at::RecordScope::USER_SCOPE) {
-            return; // only record user-defined scopes.
-          }
-          unwind::Frame frame{fn.name(), "END", 0};
-          auto r = std::make_shared<CapturedTraceback>();
-          r->recordUserDefinedFrame(frame);
-          c10::cuda::CUDACachingAllocator::recordAnnotation(r);
-        }));
+    at::addThreadLocalCallback(
+        at::RecordFunctionCallback(
+            [](const at::RecordFunction& fn)
+                -> std::unique_ptr<at::ObserverContext> {
+              unwind::Frame frame{fn.name(), "START", 0};
+              auto r = std::make_shared<CapturedTraceback>();
+              r->recordUserDefinedFrame(frame);
+              c10::cuda::CUDACachingAllocator::recordAnnotation(r);
+              return nullptr;
+            },
+            [](const at::RecordFunction& fn, at::ObserverContext* ctx_ptr) {
+              unwind::Frame frame{fn.name(), "END", 0};
+              auto r = std::make_shared<CapturedTraceback>();
+              r->recordUserDefinedFrame(frame);
+              c10::cuda::CUDACachingAllocator::recordAnnotation(r);
+            })
+            .scopes({at::RecordScope::USER_SCOPE}));
   });
 }