idk

Project.toml

@@ -5,7 +5,6 @@ version = "0.6.32"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
-Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -19,6 +18,7 @@ RewriteTools = "5969e224-3634-4c61-9f66-721b69e98b8a"
 SyntaxInterface = "b33eeca9-aacb-4496-a840-e75f1646a4fb"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
+UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
 [weakdeps]
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
@@ -34,7 +34,6 @@ TensorMarketExt = "TensorMarket"
 
 [compat]
 AbstractTrees = "0.3.4, 0.4"
-Atomix = "0.1.0"
 Compat = "3.29, 4"
 DataStructures = "0.18"
 Distributions = "0.25"
@@ -50,6 +49,7 @@ SyntaxInterface = "0.2.1"
 TOML = "1.0"
 TensorMarket = "0.2"
 UUIDs = "1.4"
+UnsafeAtomics = "0.2.1"
 julia = "1.6.7"
 
 [extras]

src/Finch.jl

@@ -19,7 +19,7 @@ using Distributions: Binomial, Normal, Poisson
 using TOML
 using UUIDs
 using Preferences
-using Atomix
+using UnsafeAtomics
 
 export @finch, @finch_program, @finch_code, @finch_kernel, value
 

src/architecture.jl

@@ -170,3 +170,29 @@ function moveto(vec::CPULocalVector, task::CPUThread)
     temp = vec.data[task.tid]
     return temp
 end
+
+function atomic_modify!(::Serial, vec, idx, op, x)
+    @inbounds begin
+        vec[idx] = op(vec[idx], x)
+    end
+end
+
+function atomic_modify!(::CPU, vec, idx, op, x)
+    Base.unsafe_modify!(pointer(vec, idx), op, x, :sequentially_consistent)
+end
+
+for T = [Bool, Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64, Int128, UInt128, Float16, Float32, Float64]
+    if T <: AbstractFloat
+        ops = [+, -]
+    else
+        ops = [+, -, *, /, %, &, |, ⊻, ⊼, max, min]
+    end
+    for op in ops
+        @eval @propagate_inbounds function atomic_modify!(::CPU, vec::Vector{$T}, idx, ::typeof($op), x::$T)
+            UnsafeAtomics.modify!(pointer(vec, idx), UnsafeAtomics.right, x, UnsafeAtomics.seq_cst)
+        end
+    end
+    @eval @propagate_inbounds function atomic_modify!(::CPU, vec::Vector{$T}, idx, ::typeof(overwrite), x::$T)
+        UnsafeAtomics.modify!(pointer(vec, idx), UnsafeAtomics.right, x, UnsafeAtomics.seq_cst)
+    end
+end

src/tensors/levels/atomic_element_levels.jl

@@ -178,7 +178,8 @@ function lower_assign(ctx, fbr::VirtualSubFiber{VirtualAtomicElementLevel}, mode
     (lvl, pos) = (fbr.lvl, fbr.pos)
     op = ctx(op)
     rhs = ctx(rhs)
-    :(Base.unsafe_modify!(pointer($(lvl.val), ($(ctx(pos)),)), $op, $rhs, :sequentially_consistent))
+    device = ctx(virtual_get_device(get_task(ctx)))
+    :(Finch.atomic_modify!($device, $(lvl.val), $(ctx(pos)), $op, $rhs))
 end
 
 function lower_assign(ctx, fbr::VirtualHollowSubFiber{VirtualAtomicElementLevel}, mode::Updater, op, rhs)
@@ -188,5 +189,6 @@ function lower_assign(ctx, fbr::VirtualHollowSubFiber{VirtualAtomicElementLevel}
     end)
     op = ctx(op)
     rhs = ctx(rhs)
-    :(Base.unsafe_modify!(pointer($(lvl.val), ($(ctx(pos)),)), $op, $rhs, :sequentially_consistent))
+    device = ctx(virtual_get_device(get_task(ctx)))
+    :(Finch.atomic_modify!($device, $(lvl.val), $(ctx(pos)), $op, $rhs))
 end