Merge pull request #2049 from CliMA/ck/dss_less_internals

Use less DataLayouts internals in DSS
CliMA · Oct 20, 2024 · 0635ff3 · 0635ff3
2 parents 78894fc + 372aea2
commit 0635ff3
Show file tree

Hide file tree

Showing 3 changed files with 139 additions and 89 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,