ext: CUDA trait

PTsolvers · Apr 29, 2024 · 1ecccb7 · 1ecccb7
1 parent 2d0030e
commit 1ecccb7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/ext/JustRelaxCUDAExt.jl b/ext/JustRelaxCUDAExt.jl
@@ -18,16 +18,13 @@ module JustRelax2D
     using MPI
 
     import JustRelax:
-        IGG, BackendTrait, CPUBackendTrait, backend, CPUBackend, Geometry, @cell
+        IGG, BackendTrait, CPUBackendTrait, CUDABackendTrait, backend, CPUBackend, Geometry, @cell
 
     @init_parallel_stencil(CUDA, Float64, 2)
 
     include("../src/common.jl")
     include("../src/stokes/Stokes2D.jl")
 
-    # add CUDA traits
-    struct CUDABackendTrait <: BackendTrait end
-
     @inline backend(::CuArray) = CUDABackendTrait()
     @inline backend(::Type{<:CuArray}) = CUDABackendTrait()
 
@@ -144,7 +141,7 @@ module JustRelax2D
     end
 
     ## Stress
-    JustRelax.JustRelax2D.tensor_invariant!(A::SymmetricTensor) = tensor_invariant!(A)
+    JustRelax.JustRelax2D.tensor_invariant!(A::JustRelax.SymmetricTensor) = tensor_invariant!(A)
 
     ## Buoyancy forces
     function JustRelax.JustRelax2D.compute_ρg!(ρg::CuArray, rheology, args)

diff --git a/src/types/traits.jl b/src/types/traits.jl
@@ -1,8 +1,8 @@
 abstract type BackendTrait end
 struct CPUBackendTrait <: BackendTrait end
 struct NonCPUBackendTrait <: BackendTrait end
-# struct CUDABackendTrait <: BackendTrait end
-# struct AMDGPUBackendTrait <: BackendTrait end
+struct CUDABackendTrait <: BackendTrait end
+struct AMDGPUBackendTrait <: BackendTrait end
 
 @inline backend(::Array) = CPUBackendTrait()
 @inline backend(::Type{<:Array}) = CPUBackendTrait()