Opt out of reduction search in some cases, fixes #288 (#289)

Author: chriselrod

src/parse/add_compute.jl

@@ -307,19 +307,17 @@ function add_compute!(
         deps = newloopdeps; reduceddeps = newreduceddeps
     end
     # @show reduction, search_tree(vparents, var) ex var vparents mpref get(ls.opdict, var, nothing) search_tree_for_ref(ls, vparents, mpref, var) # relies on cycles being forbidden
-    op = if reduction || search_tree(vparents, var)
-        add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
+    if reduction || search_tree(vparents, var)
+        return add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
     else
-        var, found = search_tree_for_ref(ls, vparents, mpref, var)
-        if found
-            add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
-        else
-            op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, vparents)
-            pushop!(ls, op, var)
+        if mpref ≢ nothing && ((length(loopdependencies(mpref)) < position) | (length(reduceddependencies(mpref)) > 0))
+            var, found = search_tree_for_ref(ls, vparents, mpref, var)
+            found && return add_reduction!(ls, var, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
         end
+        op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, vparents)
+        return pushop!(ls, op, var)
     end
     # maybe_const_compute!(ls, op, elementbytes, position)
-    op
 end
 function add_reduction!(ls::LoopSet, var::Symbol, reduceddeps, deps, vparents, reduction_ind, elementbytes, instr)
     parent = ls.opdict[var]

test/reduction_untangling.jl

@@ -0,0 +1,52 @@
+
+# issue 288
+function not_a_reduction!(A, B)
+  @turbo for j ∈ 1:size(A.re, 2)
+    jre = B.re[j]
+    jim = B.im[j]
+    for i ∈ 1:size(A.re, 1)
+      ire = B.re[i]
+      iim = B.im[i]
+      cisim = iim * jre - ire * jim
+      cisre = ire * jre + iim * jim
+      ρre_i = A.re[i,j]
+      ρim_i = A.im[i,j]
+      re_out = ρre_i * cisre - ρim_i * cisim
+      im_out = ρre_i * cisim + ρim_i * cisre
+      A.re[i,j] = re_out
+      A.im[i,j] = im_out
+    end
+  end
+  return nothing
+end
+function not_a_reduction_noturbo!(A, B)
+  @turbo for j ∈ 1:size(A.re, 2)
+    jre = B.re[j]
+    jim = B.im[j]
+    for i ∈ 1:size(A.re, 1)
+      ire = B.re[i]
+      iim = B.im[i]
+      cisim = iim * jre - ire * jim
+      cisre = ire * jre + iim * jim
+      ρre_i = A.re[i,j]
+      ρim_i = A.im[i,j]
+      re_out = ρre_i * cisre - ρim_i * cisim
+      im_out = ρre_i * cisim + ρim_i * cisre
+      A.re[i,j] = re_out
+      A.im[i,j] = im_out
+    end
+  end
+  return nothing
+end
+
+@testset "Untangle reductions" begin
+  N = 11
+  A1 = (re = rand(N,N), im = rand(N,N))
+  A2 = deepcopy(A1)
+  B = (re = rand(N), im = rand(N))
+  not_a_reduction!(A1, B)
+  not_a_reduction_noturbo!(A2, B)
+  @test A1.re ≈ A2.re
+  @test A1.im ≈ A2.im
+end
+

test/runtests.jl

@@ -56,6 +56,8 @@ const START_TIME = time()
     @time include("special.jl")
 
     @time include("multiassignments.jl")
+
+    @time include("reduction_untangling.jl")
   end
 
   @time if LOOPVECTORIZATION_TEST == "all" || LOOPVECTORIZATION_TEST == "part2"