8000 [DEBUG] REmove has CUDA · pytorch/pytorch@ccf861f · GitHub
[go: up one dir, main page]
Skip to content

Commit ccf861f

Browse files
mzzchyfacebook-github-bot
mzzchy
authored and
facebook-github-bot
committed
[DEBUG] REmove has CUDA
Differential Revision: D74549863
1 parent 0104ac0 commit ccf861f

File tree

12 files changed

+117
-7
lines changed
  • profiler

    • 12 files changed

    +117
    -7
    lines changed

    ‎aten/src/ATen/detail/MTIAHooksInterface.h

    Lines changed: 1 addition & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -126,7 +126,7 @@ struct TORCH_API MTIAHooksInterface : AcceleratorHooksInterface {
    126126
    FAIL_MTIAHOOKS_FUNC(__func__);
    127127
    }
    128128

    129-
    virtual PyObject* memorySnapshot() const {
    129+
    virtual PyObject* memorySnapshot(const std::optional<std::string>& local_path) const {
    130130
    FAIL_MTIAHOOKS_FUNC(__func__);
    131131
    return nullptr;
    132132
    }

    ‎buckbuild.bzl

    Lines changed: 2 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -183,6 +183,7 @@ THIRD_PARTY_LIBS = {
    183183
    "ruy": ["//third-party/ruy:ruy_xplat_lib", "//third_party:ruy_lib"],
    184184
    "sleef_arm": ["//third-party/sleef:sleef_arm", "//third_party:sleef_arm"],
    185185
    "typing-extensions": ["//third-party/typing-extensions:typing-extensions", "//third_party:typing-extensions"],
    186+
    "nlohmann-json": ["fbsource//third-party/nlohmann-json:nlohmann-json", "//third_party:nlohmann-json"],
    186187
    }
    187188

    188189
    def third_party(name):
    @@ -1736,6 +1737,7 @@ def define_buck_targets(
    17361737
    deps = [
    17371738
    third_party("glog"),
    17381739
    third_party("kineto"),
    1740+
    third_party("nlohmann-json"),
    17391741
    ],
    17401742
    exported_deps = [
    17411743
    ":aten_cpu",

    ‎build_variables.bzl

    Lines changed: 1 addition & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -101,6 +101,7 @@ libtorch_profiler_sources = [
    101101
    "torch/csrc/profiler/collection.cpp",
    102102
    "torch/csrc/profiler/data_flow.cpp",
    103103
    "torch/csrc/profiler/kineto_shim.cpp",
    104+
    "torch/csrc/mtia/profiler/MTIAMemoryProfiler.cpp",
    104105
    "torch/csrc/profiler/kineto_client_interface.cpp",
    105106
    "torch/csrc/profiler/orchestration/observer.cpp",
    106107
    "torch/csrc/profiler/orchestration/python_tracer.cpp",

    ‎torch/csrc/autograd/profiler_python.cpp

    Lines changed: 2 additions & 2 deletions
    Original file line numberDiff line numberDiff line change
    @@ -1162,7 +1162,7 @@ class PythonMemoryTracer final : public python_tracer::PythonMemoryTracerBase {
    11621162
    ~PythonMemoryTracer() override = default;
    11631163
    void start() override;
    11641164
    void stop() override;
    1165-
    void export_memory_history(const std::string path) override;
    1165+
    void export_memory_history(const std::string& path) override;
    11661166
    };
    11671167

    11681168
    static void toggle_memory_tracing(bool enable) {
    @@ -1197,7 +1197,7 @@ void PythonMemoryTracer::start() {
    11971197
    toggle_memory_tracing(true);
    11981198
    }
    11991199

    1200-
    void PythonMemoryTracer::export_memory_history(const std::string path) {
    1200+
    void PythonMemoryTracer::export_memory_history(const std::string& path) {
    12011201
    PyGILState_STATE gil_state = PyGILState_Ensure();
    12021202
    THPObjectPtr torch_cuda_memory_module(
    12031203
    PyImport_ImportModule("torch.cuda.memory"));

    ‎torch/csrc/mtia/Module.cpp

    Lines changed: 2 additions & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -96,7 +96,8 @@ void initModule(PyObject* module) {
    9696
    });
    9797

    9898
    m.def("_mtia_memorySnapshot", []() {
    99-
    PyObject* raw_pyobject = at::detail::getMTIAHooks().memorySnapshot();
    99+
    PyObject* raw_pyobject =
    100+
    at::detail::getMTIAHooks().memorySnapshot(std::nullopt);
    100101
    return py::reinterpret_steal<py::object>(raw_pyobject);
    101102
    });
    102103

    ‎torch/csrc/mtia/profiler/MTIAMemoryProfiler.cpp

    Lines changed: 35 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -0,0 +1,35 @@
    1+
    #include <ATen/Context.h>
    2+
    #include <ATen/detail/MTIAHooksInterface.h>
    3+
    #include <nlohmann/json.hpp>
    4+
    #include <torch/csrc/mtia/profiler/MTIAMemoryProfiler.h>
    5+
    6+
    using json = nlohmann::json;
    7+
    8+
    namespace torch::mtia {
    9+
    10+
    void MTIAMemoryProfiler::start() {
    11+
    at::detail::getMTIAHooks().recordMemoryHistory("all", "all", 150000);
    12+
    }
    13+
    14+
    void MTIAMemoryProfiler::export_memory_history(const std::string& path) {
    15+
    at::detail::getMTIAHooks().memorySnapshot(path);
    16+
    return;
    17+
    }
    18+
    19+
    void MTIAMemoryProfiler::stop() {
    20+
    at::detail::getMTIAHooks().recordMemoryHistory(std::nullopt, "all", 0);
    21+
    }
    22+
    23+
    std::unique_ptr<torch::profiler::impl::python_tracer::PythonMemoryTracerBase>
    24+
    getMemoryTracer() {
    25+
    return std::make_unique<MTIAMemoryProfiler>();
    26+
    }
    27+
    28+
    void initMemoryProfiler() {
    29+
    if (at::detail::isMTIAHooksBuilt()) {
    30+
    fprintf(stderr, "Initializing MTIA Memory Tracer\n");
    31+
    torch::profiler::impl::python_tracer::registerMemoryTracer(
    32+
    &getMemoryTracer);
    33+
    }
    34+
    }
    35+
    } // namespace torch::mtia

    ‎torch/csrc/mtia/profiler/MTIAMemoryProfiler.h

    Lines changed: 20 additions & 0 deletions
    Original file line numberDiff line numberDiff line change
    @@ -0,0 +1,20 @@
    1+
    #pragma once
    2+
    #include <torch/csrc/profiler/orchestration/python_tracer.h>
    3+
    4+
    namespace torch::mtia {
    5+
    using namespace torch::profiler::impl::python_tracer;
    6+
    7+
    void initMemoryProfiler();
    8+
    9+
    std::unique_ptr<PythonMemoryTracerBase> getMemoryTracer();
    10+
    11+
    class MTIAMemoryProfiler final : public PythonMemoryTracerBase {
    12+
    public:
    13+
    explicit MTIAMemoryProfiler() = default;
    14+
    ~MTIAMemoryProfiler() override = default;
    15+
    void start() override;
    16+
    void stop() override;
    17+
    void export_memory_history(const std::string& path) override;
    18+
    };
    19+
    20+
    } // namespace torch::mtia

    ‎torch/csrc/profiler/kineto_client_interface.cpp

    Lines changed: 4 additions & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -2,6 +2,7 @@
    22
    #include <ATen/Context.h>
    33
    #include <libkineto.h>
    44
    #include <torch/csrc/autograd/profiler_kineto.h>
    5+
    #include <torch/csrc/mtia/profiler/MTIAMemoryProfiler.h>
    56
    #include <torch/csrc/profiler/kineto_client_interface.h>
    67
    #include <chrono>
    78
    #include <thread>
    @@ -23,7 +24,9 @@ using namespace torch::autograd::profiler;
    2324

    2425
    class LibKinetoClient : public libkineto::ClientInterface {
    2526
    public:
    26-
    void init() override {}
    27+
    void init() override {
    28+
    ::torch::mtia::initMemoryProfiler();
    29+
    }
    2730

    2831
    void prepare(
    2932
    bool report_input_shapes = false,

    ‎torch/csrc/profiler/orchestration/python_tracer.cpp

    Lines changed: 1 addition & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -24,7 +24,7 @@ struct NoOpMemoryPythonTracer : public PythonMemoryTracerBase {
    2424
    ~NoOpMemoryPythonTracer() override = default;
    2525
    void start() override {}
    2626
    void stop() override {}
    27-
    void export_memory_history(const std::string path) override {}
    27+
    void export_memory_history(const std::string&) override {}
    2828
    };
    2929

    3030
    } // namespace

    ‎torch/csrc/profiler/orchestration/python_tracer.h

    Lines changed: 1 addition & 1 deletion
    Original file line numberDiff line numberDiff line change
    @@ -66,7 +66,7 @@ struct TORCH_API PythonMemoryTracerBase {
    6666

    6767
    virtual void start() = 0;
    6868
    virtual void stop() = 0;
    69-
    virtual void export_memory_history(const std::string path) = 0;
    69+
    virtual void export_memory_history(const std::string& path) = 0;
    7070
    };
    7171

    7272
    using MakeMemoryFn = std::unique_ptr<PythonMemoryTracerBase> (*)();

    0 commit comments

    Comments
     (0)
    0