xe: jit: add comments to ir codegen

Author: rjoursler

src/gpu/intel/jit/codegen/codegen.cpp

@@ -98,25 +98,28 @@ class ir_to_ngen_t : public ir_visitor_t {
             }
             expr_binding_.bind(obj.buf, rbd);
         }
-        gpu_trace() << "codegen:bind " << obj.buf << " -> "
-                    << expr_binding_.get(obj.buf);
+        host_->comment(
+                obj.line_str() + " -> " + expr_binding_.get(obj.buf).str());
         visit(obj.body);
         if (do_alloc) expr_binding_.unbind(obj.buf);
         if (use_bc_alloc) release_bank_conflict_allocation(obj);
     }
 
     void _visit(const for_t &obj) override {
+        host_->comment(obj.line_str());
         auto scope = register_scope();
         auto var_op = scope.alloc_reg_data(obj.var.type());
         bool dynamic_loop = !is_const(obj.init) || !is_const(obj.bound);
         auto init_op = eval(obj.init, scope);
         auto bound_op = eval(obj.bound, scope);
         auto step_op = eval(obj.step, scope);
 
-        host_->emov(1, var_op, init_op);
         expr_binding_.bind(obj.var, var_op);
-        gpu_trace() << "codegen:bind " << obj.var << " -> "
-                    << expr_binding_.get(obj.var);
+        host_->comment(
+                obj.var.str() + " -> " + expr_binding_.get(obj.var).str());
+
+        host_->emov(1, var_op, init_op);
+
         // For dynamic loops use standard format otherwise
         // use do-while format.
         if (dynamic_loop) {
@@ -141,9 +144,11 @@ class ir_to_ngen_t : public ir_visitor_t {
         }
 
         expr_binding_.unbind(obj.var);
+        host_->comment("end " + obj.line_str());
     }
 
     void _visit(const func_call_t &obj) override {
+        host_->comment(obj.line_str());
         auto scope = register_scope();
 
         auto &func = obj.func;
@@ -202,6 +207,7 @@ class ir_to_ngen_t : public ir_visitor_t {
 
     void _visit(const if_t &obj) override {
         gpu_assert(obj.cond.type().elems() == simd_size_);
+        host_->comment(obj.line_str());
 
         bool has_else = !obj.else_body.is_empty();
         auto scope = register_scope();
@@ -213,18 +219,20 @@ class ir_to_ngen_t : public ir_visitor_t {
                 has_else ? l_else : l_endif, l_endif);
         visit(obj.body);
         if (has_else) {
+            host_->comment("else // " + obj.line_str());
             host_->else_(simd_size_, l_endif, l_endif);
             host_->mark(l_else);
             visit(obj.else_body);
         }
         host_->mark(l_endif);
         host_->endif(simd_size_);
+        host_->comment("end " + obj.line_str());
     }
 
     void _visit(const let_t &obj) override {
         if (obj.value.is_empty()) {
             auto var_op = expr_binding_.get(obj.var);
-            gpu_trace() << "codegen:bind " << obj.var << " -> " << var_op;
+            host_->comment(obj.line_str() + " -> " + var_op.str());
             // External variable, must be already bound.
             gpu_assert(expr_binding_.is_bound(obj.var))
                     << "Variable is not defined: " << obj.var;
@@ -233,6 +241,7 @@ class ir_to_ngen_t : public ir_visitor_t {
         }
 
         auto scope = register_scope();
+        host_->comment(obj.line_str());
         if (is_const(obj.value) || is_shuffle_const(obj.value)
                 || obj.var.type() != obj.value.type()) {
             auto &var_type = obj.var.type();
@@ -247,7 +256,7 @@ class ir_to_ngen_t : public ir_visitor_t {
         }
 
         auto var_op = expr_binding_.get(obj.var);
-        gpu_trace() << "codegen:bind " << obj.var << " -> " << var_op;
+        host_->comment(obj.var.str() + " -> " + var_op.str());
 
         // At this point the scope contains allocations for temporary
         // expressions. We need to 1) query and later re-claim the allocation
@@ -279,6 +288,7 @@ class ir_to_ngen_t : public ir_visitor_t {
     }
 
     void _visit(const store_t &obj) override {
+        host_->comment(obj.line_str());
         auto scope = register_scope();
         auto buf_op = eval(obj.buf, scope);
         auto off = to_cpp<int>(obj.off);
@@ -303,6 +313,7 @@ class ir_to_ngen_t : public ir_visitor_t {
     }
 
     void _visit(const while_t &obj) override {
+        host_->comment(obj.line_str());
         auto scope = register_scope();
 
         ngen::Label loop_end_label;
@@ -314,6 +325,7 @@ class ir_to_ngen_t : public ir_visitor_t {
         visit(obj.body);
         host_->jmpi(1, loop_begin_label);
         host_->mark(loop_end_label);
+        host_->comment("end " + obj.line_str());
     }
 
 private:
@@ -1567,15 +1579,19 @@ template <typename ngen_generator_t>
 void convert_ir_to_ngen_impl(const stmt_t &body, ngen_generator_t *host,
         const walk_order_t *kernel_grid_walk_order) {
     expr_binding_t expr_binding(host->getHardware());
+    host->comment("Prologue");
     host->generate_prologue();
+
     host->bind_external_vars(body, expr_binding);
     if (kernel_grid_walk_order)
         host->bind_kernel_grid_walk_order(
                 *kernel_grid_walk_order, expr_binding);
 
+    host->comment("IR");
     ir_to_ngen_t<ngen_generator_t> visitor(host, expr_binding);
     visitor.visit(body);
 
+    host->comment("Epilogue");
     host->generate_epilogue();
 }
 

src/gpu/intel/jit/ir/core.hpp

@@ -2075,6 +2075,12 @@ class alloc_t : public stmt_impl_t {
                 : 0;
     }
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << "alloc " << buf.as<var_t>().name << "[" << size << "]";
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     expr_t buf;
@@ -2148,6 +2154,17 @@ class store_t : public stmt_impl_t {
 
     bool has_default_stride() const { return stride == default_stride; }
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << load_t::make(value.type(), buf, off, stride);
+        out << " = " << value;
+        if (!mask.is_empty()) {
+            out << ", mask = " << mask.str();
+            if (fill_mask0) out << " [FILL]";
+        }
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     static const int default_stride = -1;
@@ -2207,6 +2224,14 @@ class for_t : public stmt_impl_t {
         return ir_utils::get_hash(var, init, bound, body, step, unroll);
     }
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << "for (" << var << " = " << init << "; " << var << " < " << bound
+            << "; " << var << " += " << step << ") ";
+        if (unroll != 1) out << "[unroll: " << unroll << "] ";
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     expr_t var;
@@ -2256,6 +2281,12 @@ class if_t : public stmt_impl_t {
         return ir_utils::get_hash(cond, body, else_body);
     }
 
+    std::string line_str() const {
+        std::ostringstream oss;
+        oss << "if (" << cond << ")";
+        return oss.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     expr_t cond;
@@ -2304,6 +2335,12 @@ class let_t : public stmt_impl_t {
         return utils::rnd_up(var.type().size(), reg_allocator_t::granularity);
     };
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << var << "." << var.type() << " = " << value;
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     expr_t var;
@@ -2498,6 +2535,12 @@ class while_t : public stmt_impl_t {
 
     size_t get_hash() const override { return ir_utils::get_hash(cond, body); }
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << "while (" << cond << ")";
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     expr_t cond;
@@ -2660,6 +2703,13 @@ class func_call_t : public stmt_impl_t {
 
     size_t get_hash() const override { return ir_utils::get_hash(args, attr); }
 
+    std::string line_str() const {
+        std::ostringstream out;
+        out << func << "(" << ir_utils::make_seq_print_helper(args) << ")";
+        if (!attr.is_empty()) out << " " << attr;
+        return out.str();
+    }
+
     IR_DECLARE_TRAVERSERS()
 
     func_t func;

src/gpu/intel/jit/ir/ir.cpp

@@ -43,8 +43,7 @@ class ir_printer_t : public ir_visitor_t {
         auto grf_size = 1; // Assume all objects are grf aligned
         auto guard = mem_usage_guard(obj.register_alloc_size(grf_size));
         print_indent();
-        out_ << "alloc " << obj.buf.as<var_t>().name << "[" << obj.size
-             << "] (mem_usage: " << mem_usage_bytes_ << ")\n";
+        out_ << obj.line_str() << "(mem_usage: " << mem_usage_bytes_ << ")\n";
         visit(obj.body);
     }
 
@@ -98,16 +97,14 @@ class ir_printer_t : public ir_visitor_t {
 
     void _visit(const func_call_t &obj) override {
         print_indent();
-        out_ << obj.func << "(" << make_seq_print_helper(obj.args) << ")";
-        if (!obj.attr.is_empty()) out_ << " " << obj.attr;
-        out_ << "\n";
+        out_ << obj.line_str() << "\n";
     }
 
     void _visit(const func_impl_t &obj) override { out_ << obj.str(); }
 
     void _visit(const if_t &obj) override {
         print_indent();
-        out_ << "if (" << strip_parens(obj.cond.str()) << ") {\n";
+        out_ << obj.line_str() << " {\n";
         add_indent();
         visit(obj.body);
         remove_indent();
@@ -137,7 +134,7 @@ class ir_printer_t : public ir_visitor_t {
         int size = obj.register_alloc_size();
         auto guard = mem_usage_guard(size);
         print_indent();
-        out_ << obj.var << "." << obj.var.type() << " = " << obj.value << "\n";
+        out_ << obj.line_str() << "\n";
         visit(obj.body);
     }
 
@@ -216,13 +213,7 @@ class ir_printer_t : public ir_visitor_t {
 
     void _visit(const store_t &obj) override {
         print_indent();
-        out_ << load_t::make(obj.value.type(), obj.buf, obj.off, obj.stride);
-        out_ << " = " << obj.value;
-        if (!obj.mask.is_empty()) {
-            out_ << ", mask = " << obj.mask.str();
-            if (obj.fill_mask0) out_ << " [FILL]";
-        }
-        out_ << "\n";
+        out_ << obj.line_str() << "\n";
     }
 
     void _visit(const ternary_op_t &obj) override {
@@ -239,7 +230,7 @@ class ir_printer_t : public ir_visitor_t {
 
     void _visit(const while_t &obj) override {
         print_indent();
-        out_ << "while (" << obj.cond << ") {\n";
+        out_ << obj.line_str() << " {\n";
         add_indent();
         visit(obj.body);
         remove_indent();