Temp fix bool simd load for DTypePointer

Signed-off-by: Brian Grenier <grenierb96@gmail.com>
modularml · May 27, 2024 · e0e40ae · e0e40ae
1 parent ce285fd
commit e0e40ae
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 8 deletions.
diff --git a/stdlib/src/memory/unsafe.mojo b/stdlib/src/memory/unsafe.mojo
@@ -910,11 +910,23 @@ struct DTypePointer[
                 ]()
             return v
 
-        return (
-            self.address.offset(offset)
-            .bitcast[SIMD[type, width]]()
-            .load[alignment=alignment]()
-        )
+        # TODO: This is a temp fix for #2813
+        @parameter
+        if type.is_bool():
+            var v = SIMD[type, width]()
+
+            @parameter
+            for i in range(width):
+                v[i] = self.address.offset(int(offset) + i).load[
+                    alignment=alignment
+                ]()
+            return v
+        else:
+            return (
+                self.address.offset(offset)
+                .bitcast[SIMD[type, width]]()
+                .load[alignment=alignment]()
+            )
 
     @always_inline("nodebug")
     fn store[
@@ -960,9 +972,17 @@ struct DTypePointer[
         constrained[
             alignment > 0, "alignment must be a positive integer value"
         ]()
-        self.address.bitcast[SIMD[type, width]]().store[alignment=alignment](
-            val
-        )
+
+        @parameter
+        if type.is_bool():
+
+            @parameter
+            for i in range(width):
+                self.address.store[alignment=alignment](val[i])
+        else:
+            self.address.bitcast[SIMD[type, width]]().store[
+                alignment=alignment
+            ](val)
 
     @always_inline("nodebug")
     fn simd_nt_store[

diff --git a/stdlib/test/memory/test_memory.mojo b/stdlib/test/memory/test_memory.mojo
@@ -518,6 +518,37 @@ def test_indexing():
     assert_equal(ptr[1], 1)
 
 
+def test_wide_load():
+    var ptr = DTypePointer[DType.bool].alloc(4)
+    ptr.store(0, True)
+    ptr.store(1, True)
+    ptr.store(2, True)
+    ptr.store(3, True)
+    assert_equal(
+        ptr.load[width=4](0), SIMD[DType.bool, 4](True, True, True, True)
+    )
+
+    var ptr2 = DTypePointer[DType.int32].alloc(4)
+    ptr2.store(0, 1)
+    ptr2.store(1, 2)
+    ptr2.store(2, 3)
+    ptr2.store(3, 4)
+    assert_equal(ptr2.load[width=4](0), SIMD[DType.int32, 4](1, 2, 3, 4))
+    assert_equal(ptr2.load[width=2](2), SIMD[DType.int32, 2](3, 4))
+
+
+def test_wide_store():
+    var ptr = DTypePointer[DType.int32].alloc(4)
+    ptr.store[width=4](SIMD[DType.int32, 4](1, 2, 3, 4))
+    assert_equal(ptr.load[width=4](), SIMD[DType.int32, 4](1, 2, 3, 4))
+
+    var ptr2 = DTypePointer[DType.bool].alloc(4)
+    ptr2.store[width=4](SIMD[DType.bool, 4](True, True, True, True))
+    assert_equal(
+        ptr2.load[width=4](), SIMD[DType.bool, 4](True, True, True, True)
+    )
+
+
 def main():
     test_memcpy()
     test_memcpy_dtype()
@@ -537,3 +568,5 @@ def main():
     test_dtypepointer_gather()
     test_dtypepointer_scatter()
     test_indexing()
+    test_wide_load()
+    test_wide_store()