Remove more use of DataLayout internals in DSS

CliMA · Oct 19, 2024 · 6ca57f4 · 6ca57f4
1 parent 78894fc
commit 6ca57f4
Show file tree

Hide file tree

Showing 3 changed files with 140 additions and 98 deletions.
diff --git a/ext/cuda/data_layouts_threadblock.jl b/ext/cuda/data_layouts_threadblock.jl
@@ -170,6 +170,20 @@ end
 ##### Custom partitions
 #####
 
+##### linear partition
+@inline function linear_partition(nitems::Integer, n_max_threads::Integer)
+    threads = min(nitems, n_max_threads)
+    blocks = cld(nitems, threads)
+    return (; threads, blocks)
+end
+@inline function linear_universal_index(us::UniversalSize)
+    inds = DataLayouts.universal_size(us)
+    CI = CartesianIndices(map(x -> Base.OneTo(x), inds))
+    return CI
+end
+@inline linear_is_valid_index(i::Integer, us::UniversalSize) =
+    1 ≤ i ≤ DataLayouts.get_N(us)
+
 ##### Column-wise
 @inline function columnwise_partition(
     us::DataLayouts.UniversalSize,