From 235584eb37598d8c1935eb1d3ff788aa883df771 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Thu, 9 Nov 2023 14:54:20 -0500
Subject: [PATCH] Make CPU loops simd & ivdep

---
 src/KernelAbstractions.jl | 18 +++++++++---------
 src/extras/loopinfo.jl    |  2 ++
 src/macros.jl             |  3 +++
 3 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/src/KernelAbstractions.jl b/src/KernelAbstractions.jl
index efec0c4d1..73f1b6f69 100644
--- a/src/KernelAbstractions.jl
+++ b/src/KernelAbstractions.jl
@@ -650,6 +650,15 @@ function construct(backend::Backend, ::S, ::NDRange, xpu_name::XPUName) where {B
     return Kernel{Backend, S, NDRange, XPUName}(backend, xpu_name)
 end
 
+###
+# Extras
+# - LoopInfo
+###
+
+include("extras/extras.jl")
+
+include("reflection.jl")
+
 ###
 # Compiler
 ###
@@ -698,15 +707,6 @@ end
 __size(args::Tuple) = Tuple{args...}
 __size(i::Int) = Tuple{i}
 
-###
-# Extras
-# - LoopInfo
-###
-
-include("extras/extras.jl")
-
-include("reflection.jl")
-
 # Initialized
 
 @kernel function init_kernel(arr, f::F, ::Type{T}) where {F, T}
diff --git a/src/extras/loopinfo.jl b/src/extras/loopinfo.jl
index bb7e63850..b18a9561f 100644
--- a/src/extras/loopinfo.jl
+++ b/src/extras/loopinfo.jl
@@ -17,6 +17,8 @@ module MD
     unroll_disable() = (Symbol("llvm.loop.unroll.disable"), 1)
     unroll_enable() = (Symbol("llvm.loop.unroll.enable"), 1)
     unroll_full() = (Symbol("llvm.loop.unroll.full"), 1)
+    simd() = Symbol("julia.simdloop")
+    ivdep() = Symbol("julia.ivdep")
 end
 
 function loopinfo(expr, nodes...)
diff --git a/src/macros.jl b/src/macros.jl
index a659551d6..8766a4873 100644
--- a/src/macros.jl
+++ b/src/macros.jl
@@ -238,6 +238,8 @@ function split(stmts,
     return new_stmts
 end
 
+import .Extras: LoopInfo
+
 function emit(loop)
     idx = gensym(:I)
     for stmt in loop.indicies
@@ -285,6 +287,7 @@ function emit(loop)
                 $__validindex(__ctx__, $idx) || continue
                 $(loop.indicies...)
                 $(unblock(body))
+                $(Expr(:loopinfo, LoopInfo.MD.simd(), LoopInfo.MD.ivdep()))
             end
         end
         push!(stmts, loopexpr)