Avoid the exception branch in expand

Author: vchuravy

src/nditeration.jl

@@ -81,8 +81,40 @@ Base.length(range::NDRange) = length(blocks(range))
     CartesianIndex(nI)
 end
 
+
+"""
+    assume(cond::Bool)
+
+Assume that the condition `cond` is true. This is a hint to the compiler, possibly enabling
+it to optimize more aggressively.
+"""
+@inline assume(cond::Bool) = Base.llvmcall(("""
+        declare void @llvm.assume(i1)
+
+        define void @entry(i8) #0 {
+            %cond = icmp eq i8 %0, 1
+            call void @llvm.assume(i1 %cond)
+            ret void
+        }
+
+        attributes #0 = { alwaysinline }""", "entry"),
+    Nothing, Tuple{Bool}, cond)
+
+@inline function assume_nonzero(CI::CartesianIndices)
+    ntuple(Val(ndims(CI))) do I
+        Base.@_inline_meta
+        indices = CI.indices[I]
+        assume(indices.stop > 0)
+    end
+end
+
 Base.@propagate_inbounds function expand(ndrange::NDRange, groupidx::Integer, idx::Integer)
-    expand(ndrange, blocks(ndrange)[groupidx], workitems(ndrange)[idx])
+    # this causes a exception branch and a div
+    B = blocks(ndrange)
+    W = workitems(ndrange)
+    assume_nonzero(B)
+    assume_nonzero(W)
+    expand(ndrange, B[groupidx], workitems(ndrange)[idx])
 end
 
 Base.@propagate_inbounds function expand(ndrange::NDRange{N}, groupidx::CartesianIndex{N}, idx::Integer) where {N}