Merge pull request #31347 from JuliaLang/vc/loopinfo2

[Loopinfo] recursive conversion to metadata

src/cgutils.cpp

@@ -171,6 +171,31 @@ static Value *stringConstPtr(IRBuilder<> &irbuilder, const std::string &txt)
     }
 }
 
+// --- MDNode ---
+Metadata *to_md_tree(jl_value_t *val) {
+    if (val == jl_nothing)
+        return nullptr;
+    Metadata *MD = nullptr;
+    if (jl_is_symbol(val)) {
+        MD = MDString::get(jl_LLVMContext, jl_symbol_name((jl_sym_t*)val));
+    } else if (jl_is_bool(val)) {
+        MD = ConstantAsMetadata::get(ConstantInt::get(T_int1, jl_unbox_bool(val)));
+    } else if (jl_is_long(val)) {
+        MD = ConstantAsMetadata::get(ConstantInt::get(T_int64, jl_unbox_long(val)));
+    } else if (jl_is_tuple(val)) {
+        SmallVector<Metadata *, 8> MDs;
+        for (int f = 0, nf = jl_nfields(val); f < nf; ++f) {
+            MD = to_md_tree(jl_fieldref(val, f));
+            if (MD)
+                MDs.push_back(MD);
+        }
+        MD = MDNode::get(jl_LLVMContext, MDs);
+    } else {
+        jl_error("LLVM metadata needs to Symbol/Bool/Int or Tuple thereof");
+    }
+    return MD;
+}
+
 // --- Debug info ---
 
 static DIType *julia_type_to_di(jl_value_t *jt, DIBuilder *dbuilder, bool isboxed = false)

src/codegen.cpp

@@ -4100,34 +4100,9 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaval)
         // parse Expr(:loopinfo, "julia.simdloop", ("llvm.loop.vectorize.width", 4))
         SmallVector<Metadata *, 8> MDs;
         for (int i = 0, ie = jl_expr_nargs(ex); i < ie; ++i) {
-            jl_value_t *arg = args[i];
-            Metadata *MD;
-            if (arg == jl_nothing)
-                continue;
-            if (jl_is_symbol(arg)) {
-                MD = MDString::get(jl_LLVMContext, jl_symbol_name((jl_sym_t*)arg));
-            } else if (jl_is_tuple(arg)) {
-                // TODO: are there loopinfo with more than one field?
-                if (jl_nfields(arg) != 2)
-                    jl_error("loopinfo: only accept 2-arg tuple");
-                jl_value_t* name  = jl_fieldref(arg, 0);
-                jl_value_t* value = jl_fieldref(arg, 1);
-                if (!jl_is_symbol(name))
-                    jl_error("loopinfo: name needs to be a symbol");
-                Metadata *MDVal;
-                if(jl_is_bool(value))
-                    MDVal = ConstantAsMetadata::get(ConstantInt::get(T_int1, jl_unbox_bool(value)));
-                if(jl_is_long(value))
-                    MDVal = ConstantAsMetadata::get(ConstantInt::get(T_int64, jl_unbox_long(value)));
-                if(!MDVal)
-                    jl_error("loopinfo: value can only be a bool or a long");
-                MD = MDNode::get(jl_LLVMContext,
-                        { MDString::get(jl_LLVMContext, jl_symbol_name((jl_sym_t*)name)), MDVal });
-            }
+            Metadata *MD = to_md_tree(args[i]);
             if (MD)
                 MDs.push_back(MD);
-            else
-                jl_error("loopinfo: argument needs to be either a symbol or a tuple of type (Symbol, Union{Int, Bool}");
         }
 
         MDNode* MD = MDNode::get(jl_LLVMContext, MDs);

test/llvmpasses/Makefile

@@ -2,5 +2,7 @@ SRCDIR := $(abspath $(dir $(lastword $(MAKEFILE_LIST))))
 JULIAHOME := $(abspath $(SRCDIR)/../..)
 include $(JULIAHOME)/Make.inc
 test:
-	PATH=$(build_bindir):$(build_depsbindir):$$PATH $(build_depsbindir)/lit/lit.py -v $(SRCDIR)
+	PATH=$(build_bindir):$(build_depsbindir):$$PATH \
+	LD_LIBRARY_PATH=${build_libdir}:$$LD_LIBRARY_PATH \
+	$(build_depsbindir)/lit/lit.py -v $(SRCDIR)
 .PHONY: test

test/llvmpasses/loopinfo.jl

@@ -0,0 +1,161 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+# RUN: julia --startup-file=no %s %t && llvm-link -S %t/* -o %t/module.ll
+# RUN: cat %t/module.ll | FileCheck %s
+# RUN: cat %t/module.ll | opt -load libjulia%shlibext -LowerSIMDLoop -S - | FileCheck %s -check-prefix=LOWER
+# RUN: julia --startup-file=no %s %t -O && llvm-link -S %t/* -o %t/module.ll
+# RUN: cat %t/module.ll | FileCheck %s -check-prefix=FINAL
+using InteractiveUtils
+using Printf
+
+## Notes:
+# This script uses the `emit` function (defined at the end) to emit either
+# optimized or unoptimized LLVM IR. Each function is emitted individually and
+# `llvm-link` is used to create a single module that can be passed to opt.
+# The order in which files are emitted and linked is important since `lit` will
+# process the test cases in order.
+#
+# There are three different test prefixes defined:
+# - `CHECK`: Checks the result of codegen
+# - `LOWER`: Checks the result of -LowerSIMDLoop
+# - `FINAL`: Checks the result of running the entire pipeline
+
+# get a temporary directory
+dir = ARGS[1]
+rm(dir, force=true, recursive=true)
+mkdir(dir)
+
+# toggle between unoptimized (CHECK/LOWER) and optimized IR (FINAL).
+optimize=false
+if length(ARGS) >= 2
+    optimize = ARGS[2]=="-O"
+end
+
+# CHECK-LABEL: @julia_simdf_
+# LOWER-LABEL: @julia_simdf_
+# FINAL-LABEL: @julia_simdf_
+function simdf(X)
+    acc = zero(eltype(X))
+    @simd for x in X
+        acc += x
+# CHECK: call void @julia.loopinfo_marker(), {{.*}}, !julia.loopinfo [[LOOPINFO:![0-9]+]]
+# LOWER-NOT: llvm.mem.parallel_loop_access
+# LOWER: fadd fast double
+# LOWER-NOT: call void @julia.loopinfo_marker()
+# LOWER: br {{.*}}, !llvm.loop [[LOOPID:![0-9]+]]
+# FINAL: fadd fast <{{[0-9]+}} x double>
+    end
+    acc
+end
+
+# CHECK-LABEL: @julia_simdf2_
+# LOWER-LABEL: @julia_simdf2_
+function simdf2(X)
+    acc = zero(eltype(X))
+    @simd ivdep for x in X
+        acc += x
+# CHECK: call void @julia.loopinfo_marker(), {{.*}}, !julia.loopinfo [[LOOPINFO2:![0-9]+]]
+# LOWER: llvm.mem.parallel_loop_access
+# LOWER-NOT: call void @julia.loopinfo_marker()
+# LOWER: fadd fast double
+# LOWER: br {{.*}}, !llvm.loop [[LOOPID2:![0-9]+]]
+    end
+    acc
+end
+
+@noinline iteration(i) = (@show(i); return nothing)
+
+# CHECK-LABEL: @julia_loop_unroll
+# LOWER-LABEL: @julia_loop_unroll
+# FINAL-LABEL: @julia_loop_unroll
+@eval function loop_unroll(N)
+    for i in 1:N
+        iteration(i)
+        $(Expr(:loopinfo, (Symbol("llvm.loop.unroll.count"), 3)))
+# CHECK: call void @julia.loopinfo_marker(), {{.*}}, !julia.loopinfo [[LOOPINFO3:![0-9]+]]
+# LOWER-NOT: call void @julia.loopinfo_marker()
+# LOWER: br {{.*}}, !llvm.loop [[LOOPID3:![0-9]+]]
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL-NOT: call void @julia_iteration
+# FINAL: br
+    end
+end
+
+# Example from a GPU kernel where we want to unroll the outer loop
+# and the inner loop is a boundschecked single iteration loop.
+# The `@show` is used to bloat the loop and `FINAL-COUNT-10:` seems
+# not to be working so we duplicate the checks.
+# CHECK-LABEL: @julia_loop_unroll2
+# LOWER-LABEL: @julia_loop_unroll2
+# FINAL-LABEL: @julia_loop_unroll2
+@eval function loop_unroll2(J, I)
+    for i in 1:10
+        for j in J
+            1 <= j <= I && continue
+            @show (i,j)
+            iteration(i)
+        end
+        $(Expr(:loopinfo, (Symbol("llvm.loop.unroll.full"),)))
+# CHECK: call void @julia.loopinfo_marker(), {{.*}}, !julia.loopinfo [[LOOPINFO4:![0-9]+]]
+# LOWER-NOT: call void @julia.loopinfo_marker()
+# LOWER: br {{.*}}, !llvm.loop [[LOOPID4:![0-9]+]]
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL: call void @julia_iteration
+# FINAL-NOT: call void @julia_iteration
+    end
+end
+
+# FINAL-LABEL: @julia_notunroll
+function notunroll(J, I)
+    for i in 1:10
+        for j in J
+            1 <= j <= I && continue
+            @show (i,j)
+            iteration(i)
+# FINAL: call void @julia_iteration
+# FINAL-NOT: call void @julia_iteration
+        end
+    end
+end
+
+## Check all the MD nodes
+# CHECK: [[LOOPINFO]] = !{!"julia.simdloop"}
+# CHECK: [[LOOPINFO2]] = !{!"julia.simdloop", !"julia.ivdep"}
+# CHECK: [[LOOPINFO3]] = !{[[LOOPUNROLL:![0-9]+]]}
+# CHECK: [[LOOPUNROLL]] = !{!"llvm.loop.unroll.count", i64 3}
+# CHECK: [[LOOPINFO4]] = !{[[LOOPUNROLL2:![0-9]+]]}
+# CHECK: [[LOOPUNROLL2]] = !{!"llvm.loop.unroll.full"}
+# LOWER: [[LOOPID]] = distinct !{[[LOOPID]]}
+# LOWER: [[LOOPID2]] = distinct !{[[LOOPID2]]}
+# LOWER: [[LOOPID3]] = distinct !{[[LOOPID3]], [[LOOPUNROLL:![0-9]+]]}
+# LOWER: [[LOOPUNROLL]] = !{!"llvm.loop.unroll.count", i64 3}
+# LOWER: [[LOOPID4]] = distinct !{[[LOOPID4]], [[LOOPUNROLL2:![0-9]+]]}
+# LOWER: [[LOOPUNROLL2]] = !{!"llvm.loop.unroll.full"}
+
+# Emit LLVM IR to dir
+counter = 0
+function emit(f, tt...)
+    global counter
+    name = nameof(f)
+    open(joinpath(dir, @sprintf("%05d-%s.ll", counter, name)), "w") do io
+        code_llvm(io, f, tt, raw=true, optimize=optimize, dump_module=true, debuginfo=:none)
+    end
+    counter+=1
+end
+
+# Maintaining the order is important
+emit(simdf, Vector{Float64})
+emit(simdf2, Vector{Float64})
+emit(loop_unroll, Int64)
+emit(loop_unroll2, Int64, Int64)
+emit(notunroll, Int64, Int64)