feat(profiler): add memory flow format for profiler

GitOrigin-RevId: f0eab2639821c42af34ace1ffe556f8e2564651a

imperative/python/megengine/tools/profiler.py

+#! /usr/bin/env python3
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 import argparse
 import getopt
 import os

imperative/python/megengine/tools/svg_viewer.html

@@ -18,6 +18,8 @@
         var svg = undefined;
         var svgWidth = undefined;
         var svgHeight = undefined;
+        var resetBtn = document.getElementById('resetBtn');
+        var relocBtn = document.getElementById('relocBtn');
 
         var loadDesc = (svgElem) => {
             var mgeType = svgElem.attributes['tag:type'];
@@ -52,7 +54,7 @@
             lastColor = undefined;
         };
 
-        function recLoadSVG(svgElem) {
+        var recLoadSVG = svgElem => {
             if (svgElem.children === undefined) {
                 return;
             }
@@ -71,8 +73,37 @@
                 recLoadSVG(child);
             }
         }
+        var scaleBoard = (x, y) => {
+            var transform = 'scale(' + x + ',' + y + ')';
+            svg.setAttribute('transform', transform);
+            board.style['width'] = svgWidth * x;
+            board.style['height'] = svgHeight * y;
+        }
+        var autoScaleBoard = () => {
+            var hRangeValue = Math.sqrt(Number(hRange.value) / 10);
+            var vRangeValue = Math.sqrt(Number(vRange.value) / 10);
+            scaleBoard(Number(hRangeValue), Number(vRangeValue));
+        }
+        var zoomBoard = dScale => {
+            scale *= dScale;
+            scaleBoard(scale, scale);
+        };
 
-        function loadSVG() {
+        var reset = () => {
+            if (!svgHeight || !svgWidth) {
+                return;
+            }
+            var vScale = window.screen.availHeight / svgHeight;
+            var hScale = window.screen.availWidth / svgWidth;
+            var minScale = Math.min(hScale, vScale);
+            zoomBoard(minScale / scale);
+            window.scrollTo({
+                top: 0,
+                left: 0,
+            });
+        };
+
+        var loadSVG = () => {
             var file = fileInput.files[0];
             var reader = new FileReader();
             reader.readAsText(file, "UTF-8");
@@ -98,26 +129,10 @@
                         info.innerHTML = elemList.join('');
                     }
                 }
+                setTimeout(reset, 1);
             };
         }
-        function scaleBoard(x, y) {
-            var transform = 'scale(' + x + ',' + y + ')';
-            svg.setAttribute('transform', transform);
-            board.style['width'] = svgWidth * x;
-            board.style['height'] = svgHeight * y;
-        }
-        function autoScaleBoard() {
-            var hRangeValue = Math.sqrt(Number(hRange.value) / 10);
-            var vRangeValue = Math.sqrt(Number(vRange.value) / 10);
-            scaleBoard(Number(hRangeValue), Number(vRangeValue));
-        }
-        fileInput.onchange = loadSVG;
-        var zoomBoard = dScale => {
-            scale *= dScale;
-            scaleBoard(scale, scale);
-        };
         window.addEventListener('wheel', e => {
-            console.log(e);
             if (e.ctrlKey) {
                 e.preventDefault();
                 e.stopPropagation();
@@ -136,19 +151,39 @@
                     top: y,
                     left: x,
                 });
-                console.log('scroll', [x, y]);
+            } else if (e.altKey) {
+
             }
         }, { 'passive': false });
+        fileInput.onchange = loadSVG;
+        resetBtn.onclick = reset;
+        relocBtn.onclick = () => {
+            if (!lastElem) {
+                return;
+            }
+            y = scale * lastElem.attributes['y'].value;
+            x = scale * lastElem.attributes['x'].value;
+            window.scrollTo({
+                top: y - window.screen.availHeight/2,
+                left: x - window.screen.availWidth/2,
+                behavior: 'smooth',
+            });
+        };
     };
 </script>
 
 <body>
-    <p id="desc" style="position: fixed;bottom: 0; background-color: white;">desc</p>
-    <p id="info" style="position: fixed;top: 0; right: 0; background-color: white;">info</p>
     <p id="board"
-        style="white-space: nowrap; display: flex; justify-content: center; align-content: center; align-items: center; margin: 0;opacity: 0.7;">
+        style="white-space: nowrap; display: flex; justify-content: center; align-content: center; align-items: center; margin: 0; left: 0;">
     </p>
-    <input type='file' id='fileInput' style="position: fixed; top: 0; background-color: white;"></input>
+    <div style="display: flex; position: fixed; top: 0; left: 0; right: 0; background-color: white; flex-grow: 2;">
+        <input type='file' id='fileInput'></input>
+        <div style="flex-grow: 1;"></div>
+        <button id='resetBtn'>reset</button>
+        <button id='relocBtn'>reloc</button>
+    </div>
+    <p id="desc" style="position: fixed; bottom: 0; background-color: white;">desc</p>
+    <p id="info" style="position: fixed;top: 0; right: 0; background-color: white;">info</p>
 </body>
 
 </html>

imperative/python/megengine/tools/svg_viewer.py

+#! /usr/bin/env python3
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+import argparse
+import contextlib
+import getopt
+import http.server
+import os
+import runpy
+import sys
+import tempfile
+
+from megengine.logger import get_logger
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="megengine.tools.svg_viewer",
+        description="View SVG Graph produced bt megengine profiler",
+    )
+    parser.add_argument("-p", "--port", type=int, default=8000, help="server port")
+    parser.add_argument(
+        "-a", "--address", type=str, default="localhost", help="server address"
+    )
+    args = parser.parse_args()
+    address = args.address
+    port = args.port
+    src_filename = "svg_viewer.html"
+    dst_filename = "index.html"
+    src_path = os.path.join(os.path.dirname(__file__), src_filename)
+    url = "http://{}:{}/{}".format("localhost", port, dst_filename)
+    ssh_fwd_cmd = "ssh -L {}:localhost:{} <remote ip>".format(port, port)
+    with tempfile.TemporaryDirectory() as serve_dir:
+        dst_path = os.path.join(serve_dir, dst_filename)
+        os.symlink(src_path, dst_path)
+        os.chdir(serve_dir)
+        get_logger().info("cd to serve directory: {}, starting".format(serve_dir))
+        server = http.server.HTTPServer(
+            (address, port), http.server.SimpleHTTPRequestHandler
+        )
+        get_logger().info(
+            "server started, please visit '{}' to watch profiling result".format(url)
+        )
+        get_logger().info(
+            "if you are in remote environment, use '{}' to forward port to local".format(
+                ssh_fwd_cmd
+            )
+        )
+        try:
+            server.serve_forever()
+        except KeyboardInterrupt:
+            get_logger().info("server exiting")
+
+
+if __name__ == "__main__":
+    main()

imperative/python/megengine/utils/profiler.py

@@ -81,6 +81,7 @@ class Profiler(ContextDecorator):
         for opt, optval in Profiler.valid_options.items():
             self._options[opt] = int(kwargs.pop(opt, optval))
         self._pid = "<PID>"
+        self._dump_callback = None
 
     @property
     def path(self):

imperative/src/impl/interpreter/interpreter_impl.cpp

@@ -48,7 +48,11 @@ namespace mgb {
     using namespace profiler;
 }
 
-#ifdef __GNUG__
+#if defined(_WIN32) || defined(_WIN64)
+#define SYMBOL_EXPORT __declspec(dllexport)
+#else
+#define SYMBOL_EXPORT __attribute__((visibility("default")))
+#endif
 
 namespace mgb {
 
@@ -62,17 +66,17 @@ namespace mgb {
  *     mgb::imperative_log_profile("MY MESSAGE");
  *
  **/
-__attribute__((visibility("default")))
+SYMBOL_EXPORT
 void imperative_log_profile_begin(const char* message) {
     RECORD_EVENT(CustomEvent, std::string{message});
 }
 
-__attribute__((visibility("default")))
+SYMBOL_EXPORT
 void imperative_log_profile_end(const char* message) {
     RECORD_EVENT(CustomFinishEvent, std::string{message});
 }
 
-__attribute__((visibility("default")))
+SYMBOL_EXPORT
 void imperative_log_profile(const char* message){
     imperative_log_profile_begin(message);
     imperative_log_profile_end(message);
@@ -80,8 +84,6 @@ void imperative_log_profile(const char* message){
 
 }
 
-#endif
-
 std::thread::id ChannelImpl::get_worker_tid() {
     return m_worker_state.tid;
 }

imperative/src/impl/profiler.cpp

@@ -73,6 +73,8 @@ void Profiler::dump_profile(std::string basename, std::string format, results_t
     auto thread_dict = get_thread_dict();
     if (format == "chrome_timeline.json") {
         profiler::dump_chrome_timeline(basename, options, thread_dict, results);
+    } else if (format == "memory_flow.svg") {
+        profiler::dump_memory_flow(basename, options, thread_dict, results);
     } else {
         mgb_log_error("unsupported profiling format %s", format.c_str());
     }

imperative/src/impl/profiler/chrome_timeline.cpp

@@ -161,12 +161,10 @@ public:
             event_list->add(event.to_json());
         }
         (*result)["traceEvents"] = event_list;
-        //(*result)["localTime"] = json::String::make(std::to_string((double)m_local_time/1e3));
         return result;
     }
 private:
     std::vector<ChromeTraceEvent> m_content;
-    uint64_t m_local_time;
 };
 
 
@@ -408,9 +406,7 @@ void dump_chrome_timeline(std::string filename, Profiler::options_t options, Pro
     });
 
     HANDLE_EVENT(TensorWaitPropEvent, {
-        auto& tensor_state = state.tensors[event.tensor_id];
         NEW_HOST("TensorWaitProp", 'B');
-                //.args(TENSOR_PROPS);
         if (event.prop == TensorProp::HostValue) {
             INC_COUNTER(wait_value_count, 1);
         } else if (event.prop == TensorProp::Shape) {
@@ -433,7 +429,6 @@ void dump_chrome_timeline(std::string filename, Profiler::options_t options, Pro
     });
 
     HANDLE_EVENT(TensorNotifyPropEvent, {
-        auto& tensor_state = state.tensors[event.tensor_id];
         NEW_HOST(ssprintf("%d", pid), 's')
                 .id(event.tensor_id)
                 .cat("TensorProp")
@@ -471,9 +466,7 @@ void dump_chrome_timeline(std::string filename, Profiler::options_t options, Pro
     });
 
     HANDLE_EVENT(TensorCommandEvent, {
-        auto& tensor_state = state.tensors[event.tensor_id];
         NEW_HOST(ssprintf("%s %zu", TENSOR_COMMAND_KIND, event.tensor_id), 'B');
-            //.args(TENSOR_PROPS);
     });
 
     HANDLE_EVENT(TensorCommandFinishEvent, {

imperative/src/impl/profiler/formats.h

@@ -19,4 +19,6 @@ namespace mgb::imperative::profiler {
 
 void dump_chrome_timeline(std::string filename, Profiler::options_t options, Profiler::thread_dict_t thread_dict, Profiler::results_t results);
 
+void dump_memory_flow(std::string filename, Profiler::options_t options, Profiler::thread_dict_t thread_dict, Profiler::results_t results);
+
 }

imperative/src/impl/profiler/memory_chunk.cpp

+#include <map>
+#include <vector>
+#include <array>
+
+#include "megbrain/imperative/utils/to_string.h"
+#include "megbrain/utils/debug.h"
+
+#include "./formats.h"
+#include "./states.h"
+
+#include "./events.h"
+
+namespace mgb::imperative::profiler {
+
+class XMLWriter {
+private:
+    std::vector<std::vector<std::string>> elements;
+public:
+    struct ElementGuard {
+        XMLWriter* writer;
+        std::string name;
+        std::vector<std::pair<std::string, std::string>> attrs;
+
+        template <typename T>
+        ElementGuard& attr(std::string key, T&& value) {
+            attrs.push_back({key, mgb::imperative::to_string(value)});
+            return *this;
+        }
+
+        std::string to_string_start() const {
+            std::string builder;
+            builder.append(ssprintf("<%s",
+                        name.c_str()));
+            for (auto&& [k, v]: attrs) {
+                builder.append(ssprintf(" %s=\"%s\"", k.c_str(), v.c_str()));
+            }
+            builder.append(">\n");
+            return builder;
+        }
+
+        std::string to_string_end() const {
+            return ssprintf("</%s>\n", name.c_str());
+        }
+
+        ElementGuard(XMLWriter* writer, std::string name): writer{writer}, name{name} {
+            writer->elements.emplace_back();
+        }
+
+        ~ElementGuard() {
+            auto children = std::move(writer->elements.back());
+            writer->elements.pop_back();
+            std::string builder;
+            builder.append(to_string_start());
+            for (auto&& child: children) {
+                builder.append(child);
+            }
+            builder.append(to_string_end());
+            writer->elements.back().push_back(builder);
+        }
+    };
+    XMLWriter() {
+        elements.emplace_back().push_back("<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n");
+    }
+    ElementGuard element(std::string tag) {
+        return ElementGuard{this, tag};
+    }
+    void text(std::string text) {
+        elements.back().push_back(text);
+    }
+    void doctype(std::string element, std::string dtd, std::vector<std::string> args) {
+        std::string builder = ssprintf("<!DOCTYPE %s %s", element.c_str(), dtd.c_str());
+        for (auto&& arg: args) {
+            builder.append(ssprintf(" %s", arg.c_str()));
+        }
+        builder.append(">\n");
+        elements.back().push_back(builder);
+    }
+    std::string to_string() const {
+        mgb_assert(elements.size() == 1 && elements[0].size() >= 1);
+        std::string builder;
+        for (auto&& element: elements[0]) {
+            builder.append(element);
+        }
+        return builder;
+    }
+};
+
+struct MemoryChunk {
+    std::array<uintptr_t, 2> address;
+    std::string name;
+    TensorLayout layout;
+    std::array<uint64_t, 2> time;
+
+    bool empty() const {
+        return address[1] - address[0] == 0;
+    }
+};
+
+struct MemoryFlow {
+    std::unordered_map<uint64_t, MemoryChunk> chunks;
+
+    std::pair<uintptr_t, uintptr_t> address_range() const {
+        auto addr_begin = std::numeric_limits<uintptr_t>::max();
+        auto addr_end = std::numeric_limits<uintptr_t>::min();
+        for(auto&& [id, chunk]: chunks) {
+            if (chunk.empty()) continue;
+            addr_begin = std::min(addr_begin, chunk.address[0]);
+            addr_end = std::max(addr_end, chunk.address[1]);
+        }
+        return {addr_begin, addr_end};
+    }
+
+    std::pair<uint64_t, uint64_t> time_range() const {
+        auto time_begin = std::numeric_limits<uint64_t>::max();
+        auto time_end = std::numeric_limits<uint64_t>::min();
+        for(auto&& [id, chunk]: chunks) {
+            if (chunk.empty()) continue;
+            time_begin = std::min(time_begin, chunk.time[0]);
+            time_end = std::max(time_end, chunk.time[1]);
+        }
+        return {time_begin, time_end};
+    }
+
+    std::shared_ptr<json::Array> to_json() const {
+        auto results = json::Array::make();
+        for(auto&& [id, chunk]: chunks) {
+            if (chunk.empty()) continue;
+            auto address = json::Array::make();
+            auto time = json::Array::make();
+            address->add(json::String::make(std::to_string(chunk.address[0])));
+            address->add(json::String::make(std::to_string(chunk.address[1])));
+            time->add(json::String::make(std::to_string(chunk.time[0])));
+            time->add(json::String::make(std::to_string(chunk.time[1])));
+            results->add(json::Object::make({
+                {"address", address},
+                {"name", json::String::make(chunk.name)},
+                {"layout", json::String::make(chunk.layout.to_string())},
+                {"time", time}
+            }));
+        }
+        return results;
+    }
+
+    XMLWriter to_svg() const {
+        XMLWriter writer;
+        auto&& [addr_begin, addr_end] = address_range();
+        auto&& [time_begin, time_end] = time_range();
+        writer.doctype("svg", "PUBLIC", {
+                "\"-//W3C//DTD SVG 1.1//EN\"",
+                "\"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\""
+        });
+        auto svg = writer.element("svg");
+        svg.attr("xmlns", std::string{"http://www.w3.org/2000/svg"});
+        svg.attr("xmlns:tag", std::string{"https://megengine.org.cn"});
+        double time_scale = 1e5;
+        double addr_scale = 1e6;
+        svg.attr("width", (time_end-time_begin)/time_scale);
+        svg.attr("height", (addr_end-addr_begin)/addr_scale);
+        {
+            auto rect = writer.element("rect");
+            rect.attr("x", 0);
+            rect.attr("y", 0);
+            rect.attr("width", (time_end-time_begin)/time_scale);
+            rect.attr("height", (addr_end-addr_begin)/addr_scale);
+            rect.attr("fill", std::string{"blue"});
+        }
+        double us = 1e3, ms = 1e6;
+        std::map<double, std::string> time2color = {
+            {0 * us, "#DDDDDD"},
+            {100 * us, "#CCCCCC"},
+            {1 * ms, "#BBBBBB"},
+            {10 * ms, "#AAAAAA"},
+            {100 * ms, "#999999"},
+            {1000 * ms, "#888888"},
+            {std::numeric_limits<double>::infinity(), "#555555"},
+        };
+        auto time2str = [](uint64_t ns){
+            using pair_t = std::pair<uint64_t, const char*>;
+            static pair_t units[] = {
+                {1, "ns "},
+                {1e3, "us "},
+                {1e6, "ms "},
+                {1e9, "s "},
+            };
+            std::string builder;
+            auto comparator = [](const pair_t& lhs, const pair_t& rhs) {
+                return lhs.first < rhs.first;
+            };
+            while (ns > 0) {
+                auto iter = std::upper_bound(std::begin(units), std::end(units), std::make_pair(ns, ""), comparator) - 1;
+                builder += std::to_string(ns / iter->first) + iter->second;
+                ns = ns % iter->first;
+            }
+            return builder;
+        };
+        auto size2str = [](size_t sz){
+            using pair_t = std::pair<size_t, const char*>;
+            static pair_t units[] = {
+                {1, "B "},
+                {1024, "KB "},
+                {1024*1024, "MB "},
+                {1024*1024*1024, "GB "},
+            };
+            std::string builder;
+            auto comparator = [](const pair_t& lhs, const pair_t& rhs) {
+                return lhs.first < rhs.first;
+            };
+            while (sz > 0) {
+                auto iter = std::upper_bound(std::begin(units), std::end(units), std::make_pair(sz, ""), comparator) - 1;
+                builder += std::to_string(sz / iter->first) + iter->second;
+                sz = sz % iter->first;
+            }
+            return builder;
+        };
+        for (auto&& [id, chunk]: chunks) {
+            if (chunk.empty()) continue;
+            double left = (chunk.time[0]-time_begin)/time_scale;
+            double right = (chunk.time[1]-time_begin)/time_scale;
+            double top = (chunk.address[0]-addr_begin)/addr_scale;
+            double bottom = (chunk.address[1]-addr_begin)/addr_scale;
+            double duration = chunk.time[1] - chunk.time[0];
+            {
+                auto rect = writer.element("rect");
+                rect.attr("x", left);
+                rect.attr("y", top);
+                rect.attr("height", bottom - top);
+                rect.attr("width", right - left);
+                rect.attr("fill", time2color.lower_bound(duration)->second);
+                auto mge_attr = [&](const char* name, auto&& value) {
+                    rect.attr(ssprintf("tag:%s", name), value);
+                };
+                mge_attr("type", std::string("tensor"));
+                mge_attr("name", chunk.name);
+                mge_attr("address", ssprintf("%p", reinterpret_cast<void*>(chunk.address[0])));
+                mge_attr("size", size2str(chunk.address[1] - chunk.address[0]));
+                mge_attr("layout", chunk.layout.to_string());
+                mge_attr("produced", time2str(chunk.time[0]));
+                mge_attr("erased", time2str(chunk.time[1]));
+                mge_attr("duration", time2str(chunk.time[1] - chunk.time[0]));
+            }
+        }
+        return writer;
+    }
+};
+
+void dump_memory_flow(std::string filename, Profiler::options_t options, Profiler::thread_dict_t thread_dict, Profiler::results_t results) {
+    MemoryFlow flow;
+
+    ProfileDataCollector collector;
+    ProfileState state;
+#define HANDLE_EVENT(type, ...) \
+    collector.handle<type>([&](uint64_t id, std::thread::id tid, uint64_t time, type event) __VA_ARGS__ );
+
+    HANDLE_EVENT(TensorDeclareEvent, {
+        auto& tensor_state = state.tensors[event.tensor_id] = {};
+        tensor_state.id = event.tensor_id;
+        tensor_state.name = event.name;
+    });
+
+    HANDLE_EVENT(TensorProduceEvent, {
+        auto& tensor_state = state.tensors[event.tensor_id];
+        tensor_state.device = event.device;
+        tensor_state.layout = event.layout;
+        tensor_state.produced = time;
+        state.tensors_by_size.insert({tensor_state.id, tensor_state.size_in_bytes()});
+        state.tensors_by_produced.insert({tensor_state.id, tensor_state.produced});
+        auto& chunk = flow.chunks[event.tensor_id];
+        uintptr_t address = reinterpret_cast<uintptr_t>(event.ptr);
+        auto span = event.layout.span();
+        auto dtype = event.layout.dtype;
+        // assume dtype is not lowbit
+        if (!address) {
+            chunk.address = {0, 0};
+        } else {
+            chunk.address = {address+span.low_elem*dtype.size(), address+span.high_elem*dtype.size()};
+        }
+        chunk.layout = tensor_state.layout;
+        chunk.time[0] = time;
+        chunk.name = tensor_state.name;
+    });
+
+    HANDLE_EVENT(TensorReleaseEvent, {
+        auto& tensor_state = state.tensors[event.tensor_id];
+        state.tensors_by_size.erase({tensor_state.id, tensor_state.size_in_bytes()});
+        state.tensors_by_produced.erase({tensor_state.id, tensor_state.produced});
+        auto& chunk = flow.chunks[event.tensor_id];
+        chunk.time[1] = time;
+    });
+
+    HANDLE_EVENT(ScopeEvent, {
+        state.threads[tid].scope_stack.push_back(event.name);
+    });
+
+    HANDLE_EVENT(ScopeFinishEvent, {
+        mgb_assert(state.threads[tid].scope_stack.back() == event.name);
+        state.threads[tid].scope_stack.pop_back();
+    });
+
+    for (auto&& result: results) {
+        collector(result.second.id, result.first, result.second.time, result.second.data);
+    }
+
+    debug::write_to_file(filename.c_str(), flow.to_svg().to_string());
+}
+
+}