Profiler: Add counters and recompute statistics

src/IR.cpp

@@ -818,7 +818,10 @@ constexpr const char *intrinsic_op_names[] = {
     "count_leading_zeros",
     "count_trailing_zeros",
     "debug_to_file",
+    "declare_allocation",
+    "declare_box_required_at_root",
     "declare_box_touched",
+    "declare_stage",
     "div_round_to_zero",
     "dynamic_shuffle",
     "extract_bits",

src/IR.h

@@ -725,8 +725,23 @@ struct Call : public ExprNode<Call> {
         count_leading_zeros,
         count_trailing_zeros,
         debug_to_file,
+        // Emitted (when profiling) for a device-only allocation that has no
+        // host-side Allocate node: the buffer lives solely on the device, so
+        // InjectHostDevBufferCopies elides its host allocation. Lets the
+        // profiler, which tracks memory at host Allocate nodes, still see the
+        // allocation and bill its size to the Func.
+        // Args: (StringImm Func name, size in bytes, IntImm memory type).
+        declare_allocation,
+        // Declares that region required of a particular Func at this
+        // scope. Injected by ScheduleFunctions and used by the profiler.
+        declare_box_required_at_root,
         // Declares that a box region of an allocation has been touched (used by bounds inference)
         declare_box_touched,
+        // Declares that the following stmt computes a particular stage of
+        // a particular Func. Used by the profiler to bill points computed
+        // in the pure def separately from points computed in update defs.
+        // Args: (Variable<Handle> handle for the func, Int<32> stage_idx).
+        declare_stage,
         div_round_to_zero,
         // A shuffle operation with runtime-varying indices.
         dynamic_shuffle,

src/InjectHostDevBufferCopies.cpp

@@ -6,6 +6,7 @@
 #include "IRMutator.h"
 #include "IROperator.h"
 #include "IRPrinter.h"
+#include "Simplify.h"
 #include "Substitute.h"
 
 #include <map>
@@ -490,6 +491,12 @@ class InjectBufferCopies : public IRMutator {
 protected:
     using IRMutator::visit;
 
+    // Whether the pipeline is being profiled, so we know whether to emit
+    // allocation markers for device-only buffers whose host allocation we
+    // strip below (the profiler tracks memory at the host Allocate, which
+    // no longer carries a size for these).
+    bool profiling;
+
     // Inject the registration of a device destructor just after the
     // .buffer symbol is defined (which is safely before the first
     // device_malloc).
@@ -658,6 +665,25 @@ class InjectBufferCopies : public IRMutator {
                 // references to it (e.g. the one in the make_buffer
                 // call) with NULL.
                 body = substitute(op->name, reinterpret(Handle(), make_zero(UInt(64))), body);
+                if (profiling) {
+                    // The storage still exists on the device, but with no
+                    // host Allocate size the profiler can't see it. Emit a
+                    // declare_allocation marker carrying the device
+                    // allocation's byte size so it gets billed to this Func.
+                    // A 0-dimensional allocation (empty extents) holds one
+                    // element, so start the product at the element size.
+                    Expr size_bytes = make_const(UInt(64), op->type.bytes());
+                    for (const Expr &extent : op->extents) {
+                        size_bytes *= cast<uint64_t>(extent);
+                    }
+                    size_bytes = simplify(size_bytes);
+                    Expr marker = Call::make(Int(32), Call::declare_allocation,
+                                             {Expr(op->name),
+                                              size_bytes,
+                                              make_const(Int(32), (int)op->memory_type)},
+                                             Call::Intrinsic);
+                    body = Block::make(Evaluate::make(marker), body);
+                }
             }
 
             return op->with(op->extents, condition, body);
@@ -674,6 +700,11 @@ class InjectBufferCopies : public IRMutator {
             return IRMutator::visit(op);
         }
     }
+
+public:
+    InjectBufferCopies(bool profiling)
+        : profiling(profiling) {
+    }
 };
 
 // Find the site in the IR where we want to inject the copies/dirty
@@ -794,7 +825,8 @@ Stmt inject_host_dev_buffer_copies(Stmt s, const Target &t) {
     }
 
     // Handle internal allocations
-    s = InjectBufferCopies()(s);
+    bool profiling = t.has_feature(Target::Profile) || t.has_feature(Target::ProfileByTimer);
+    s = InjectBufferCopies(profiling)(s);
 
     // Handle inputs and outputs
     FindOutermostProduce outermost;

src/Lower.cpp

@@ -404,12 +404,6 @@ void lower_impl(const vector<Function> &output_funcs,
     s = bound_small_allocations(s);
     log("Lowering after bounding small allocations:", s);
 
-    if (t.has_feature(Target::Profile) || t.has_feature(Target::ProfileByTimer)) {
-        debug(1) << "Injecting profiling...\n";
-        s = inject_profiling(s, pipeline_name, env);
-        log("Lowering after injecting profiling:", s);
-    }
-
     if (t.has_feature(Target::CUDA)) {
         debug(1) << "Injecting warp shuffles...\n";
         s = lower_warp_shuffles(s, t);
@@ -440,6 +434,13 @@ void lower_impl(const vector<Function> &output_funcs,
     s = hoist_loop_invariant_if_statements(s);
     log("Lowering after removing dead allocations and hoisting loop invariants:", s);
 
+    if (t.has_feature(Target::Profile) || t.has_feature(Target::ProfileByTimer)) {
+        debug(1) << "Injecting profiling...\n";
+        s = inject_profiling(s, pipeline_name, env, t);
+        s = simplify(s);
+        log("Lowering after injecting profiling:", s);
+    }
+
     debug(1) << "Finding intrinsics...\n";
     // Must be run after the last simplification, because it turns
     // divisions into shifts, which the simplifier reverses.