record function reproducer

add.py

import torch
import sys

device = torch.device('cpu')
left = torch.zeros(100, device=device, requires_grad=True)
right = torch.zeros(100, device=device, requires_grad=True)
grad = torch.zeros(100, device=device)

for _ in range(10):
    output = torch.add(left, right)
    output.backward(grad)

wrapper.cc

#include <torch/extension.h>
#include <iostream>

int driver_register() {
  at::addGlobalCallback(
    at::RecordFunctionCallback(
      [](const at::RecordFunction& fn)
      -> std::unique_ptr<at::ObserverContext> {
        std::cout << fn.forwardThreadId() << std::endl;
        return nullptr;
      },
      [](const at::RecordFunction& fn, at::ObserverContext* ctx_ptr) {
        return;
      })
    .needsInputs(false)   // TODO(Keren): monitor inputs if needed?
    .needsOutputs(false)  // TODO(Keren): monitor outputs if needed?
    .scopes({}));
  return 0;
}

int _ret = driver_register();

• Build

# whereis "include/python"
export TORCH_DIR=<path/to/pytorch/dir>
export PYTHON_DIR=<path/to/python/dir>
g++ -fPIC -std=c++17 -O3 -g -I$TORCH_DIR/include -I$TORCH_DIR/include/torch/csrc/api/include -I$PYTHON_DIR/include/python3.10 -shared -L$TORCH_DIR/lib -Wl,-rpath=$TORCH_DIR/lib wrapper.cc -o wrapper.so -lc10 -lc10_cuda -ltorch_cpu -ltorch

• Run

LD_PRELOAD=`pwd`/wrapper.so python add.py

Just to consolidate the discussion here, @Jokeren said I always get fwd_thread_id = 18446744073709551615, which should be 0 or 1