[******][GPU] Make unsafe_pointer GPU agnostic

MODULAR_ORIG_COMMIT_REV_ID: 85797bf247f55c80748fc56d468a9d8c2e0e2767
modular · Nov 14, 2024 · 076c9b9 · sibarras · Nov 14, 2024 · 076c9b9
1 parent 5db93d2
commit 076c9b9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/stdlib/src/memory/unsafe_pointer.mojo b/stdlib/src/memory/unsafe_pointer.mojo
@@ -19,7 +19,7 @@ from memory import UnsafePointer
 ```
 """
 
-from sys import alignof, sizeof, is_nvidia_gpu
+from sys import alignof, sizeof, is_nvidia_gpu, is_gpu
 from sys.intrinsics import (
     _mlirtype_is_eq,
     _type_is_eq,
@@ -40,7 +40,7 @@ from memory.memory import _free, _malloc
 
 @always_inline
 fn _default_alignment[type: AnyType]() -> Int:
-    return alignof[type]() if is_nvidia_gpu() else 1
+    return alignof[type]() if is_gpu() else 1
 
 
 @always_inline
@@ -842,7 +842,7 @@ struct UnsafePointer[
         type: DType, //,
         *,
         width: Int = 1,
-        alignment: Int = alignof[SIMD[type, width]]() if is_nvidia_gpu() else 1,
+        alignment: Int = _default_alignment[type, width](),
     ](
         self: UnsafePointer[Scalar[type], *_, **_],
         offset: SIMD[_, width],

diff --git a/stdlib/src/sys/__init__.mojo b/stdlib/src/sys/__init__.mojo
@@ -49,6 +49,7 @@ from .info import (
     warpsize,
     sizeof,
     is_nvidia_gpu,
+    is_gpu,
 )
 from .intrinsics import (
     PrefetchCache,