Fix check for u1loop axis in unrolled_curly

Author: chriselrod

src/codegen/lower_memory_common.jl

@@ -204,7 +204,7 @@ function unrolled_curly(op::Operation, u₁::Int, u₁loop::Loop, vloop::Loop, m
                 @assert AV == -1
                 AV = n
             end
-            if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp))
+            if (u₁loopsym === CONSTANTZEROINDEX) ? (CONSTANTZEROINDEX ∈ loopdependencies(opp)) : (isu₁unrolled(opp) || (ind === u₁loopsym))
                 @assert AU == -1
                 AU = n
             end
@@ -214,6 +214,11 @@ function unrolled_curly(op::Operation, u₁::Int, u₁loop::Loop, vloop::Loop, m
     #     return mem_offset_u(op, td, inds_calc_by_ptr_offset, true)
     # end
     # ind = mem_offset_u(op, td, inds_calc_by_ptr_offset, false)
+    # if AU == -1
+    #     @show op indices u₁loopsym vloopsym
+    # end
+    AU == -1 && throw(LoopError("Failed to find $(u₁loopsym) in args of $(repr(op))."))
+    # @assert AU ≠ -1
     vecnotunrolled = AU != AV
     conditional_memory_op = isconditionalmemop(op)
     if mask || conditional_memory_op

src/reconstruct_loopset.jl

@@ -630,7 +630,7 @@ Execute an `@avx` block. The block's code is represented via the arguments:
 @generated function _avx_!(
     ::Val{UNROLL}, ::Val{OPS}, ::Val{ARF}, ::Val{AM}, ::Val{LPSYM}, var"#lv#tuple#args#"::Tuple{LB,V}
 ) where {UNROLL, OPS, ARF, AM, LPSYM, LB, V}
-    1 + 1 # Irrelevant line you can comment out/in to force recompilation...
+    # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
     ls = _avx_loopset(OPS, ARF, AM, LPSYM, LB.parameters, V.parameters, UNROLL)
     # return @show avx_body(ls, UNROLL)
     if last(UNROLL) > 1