add Polyester.jl for naive multithreading of with_signature functions

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.1"
 manifest_format = "2.0"
-project_hash = "56b882753719bd60fcc27014871a9fb61f8ede1b"
+project_hash = "171fb4ccce9a2594888302a8eb8bc0fcb4ab7a3a"
 
 [[deps.Accessors]]
 deps = ["CompositionsBase", "ConstructionBase", "Dates", "InverseFunctions", "LinearAlgebra", "MacroTools", "Test"]
@@ -495,6 +495,12 @@ deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "
 uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 version = "1.10.0"
 
+[[deps.Polyester]]
+deps = ["ArrayInterface", "BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "ManualMemory", "PolyesterWeave", "Requires", "Static", "StaticArrayInterface", "StrideArraysCore", "ThreadingUtilities"]
+git-tree-sha1 = "b3e2bae88cf07baf0a051fe09666b8ef97aefe93"
+uuid = "f517fe37-dbe3-4b94-8317-1923a5111588"
+version = "0.7.14"
+
 [[deps.PolyesterWeave]]
 deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"]
 git-tree-sha1 = "240d7170f5ffdb285f9427b92333c3463bf65bf6"
@@ -668,6 +674,12 @@ git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d"
 uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 version = "1.4.2"
 
+[[deps.StrideArraysCore]]
+deps = ["ArrayInterface", "CloseOpenIntervals", "IfElse", "LayoutPointers", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface", "ThreadingUtilities"]
+git-tree-sha1 = "25349bf8f63aa36acbff5e3550a86e9f5b0ef682"
+uuid = "7792a7ef-975c-4747-a70f-980b88e8d1da"
+version = "0.5.6"
+
 [[deps.StringManipulation]]
 deps = ["PrecompileTools"]
 git-tree-sha1 = "a04cabe79c5f01f4d723cc6704070ada0b9d46d5"

Project.toml

@@ -14,6 +14,7 @@ LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

src/with_signature.jl

@@ -153,6 +153,12 @@ function parse_options(ex=_empty_opts)
       end
       options[:debug] = rhs
       continue
+    elseif lhs === :polyester
+      if !(rhs isa Bool)
+        error("@with_signature: invalid value '$rhs' for option 'polyester', expected a Bool")
+      end
+      options[:polyester] = rhs
+      continue
     end
     error("@with_signature: unknown option '$a' found in '$ex', skipping ...")
   end
@@ -162,9 +168,16 @@ function parse_options(ex=_empty_opts)
   get!(options, :inbounds, true)
   get!(options, :simd, true)
   get!(options, :debug, false)
+  get!(options, :polyester, false)
+  if options[:polyester]
+    # TODO Does @inbounds, @simd compose with @batch?
+    options[:inbounds]  = false
+    options[:simd]      = false
+  end
   if options[:debug]
-    options[:inbounds] = false
-    options[:simd]     = false
+    options[:inbounds]  = false
+    options[:simd]      = false
+    options[:polyester] = false
   end
   return options
 end
@@ -341,8 +354,7 @@ function make_with_signature_methods(expr, options=(;))
   end
 
   # make gen'ed symbols for construction of the rest of the function body
-  ii, vi, bi, vidxs, bidxs =
-    gensym.((:ii,:vi,:bi,:vidxs,:bidxs))
+  ii, vi, bi, vidxs, bidxs = gensym.((:ii,:vi,:bi,:vidxs,:bidxs))
 
   # for each variable in @accepts
   # 1) resolve the prefix,
@@ -455,6 +467,9 @@ function make_with_signature_methods(expr, options=(;))
   if options[:inbounds]
     loop_ex = Expr(:macrocall, Symbol("@inbounds"), nothing, loop_ex)
   end
+  if options[:polyester]
+    loop_ex = Expr(:macrocall, Symbol("@batch"), nothing, loop_ex)
+  end
 
   # combine loop with accepts, returns array unpacking
   new_func = deepcopy(ex)