[Bug] [vulkan] Fix data type alignment for arguments and return values (

taichi-dev#3571)
sjwsl · Nov 21, 2021 · 787271b · 787271b
1 parent 813ddae
commit 787271b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/taichi/backends/vulkan/kernel_utils.cpp b/taichi/backends/vulkan/kernel_utils.cpp
@@ -95,6 +95,9 @@ KernelContextAttributes::KernelContextAttributes(const Kernel &kernel)
     // Put scalar args in the memory first
     for (int i : scalar_indices) {
       auto &attribs = (*vec)[i];
+      const size_t dt_bytes = vk_data_type_size(attribs.dt);
+      // Align bytes to the nearest multiple of dt_bytes
+      bytes = (bytes + dt_bytes - 1) / dt_bytes * dt_bytes;
       attribs.offset_in_mem = bytes;
       bytes += attribs.stride;
       TI_TRACE("  at={} scalar offset_in_mem={} stride={}", i,
@@ -103,6 +106,8 @@ KernelContextAttributes::KernelContextAttributes(const Kernel &kernel)
     // Then the array args
     for (int i : array_indices) {
       auto &attribs = (*vec)[i];
+      const size_t dt_bytes = vk_data_type_size(attribs.dt);
+      bytes = (bytes + dt_bytes - 1) / dt_bytes * dt_bytes;
       attribs.offset_in_mem = bytes;
       bytes += attribs.stride;
       TI_TRACE("  at={} array offset_in_mem={} stride={}", i,

diff --git a/tests/python/test_arg_alignment.py b/tests/python/test_arg_alignment.py
@@ -1,7 +1,7 @@
 import taichi as ti
 
 
-@ti.test(exclude=[ti.opengl, ti.vulkan])
+@ti.test(exclude=[ti.opengl])
 def test_ret_write():
     @ti.kernel
     def func(a: ti.i16) -> ti.f32:
@@ -10,7 +10,7 @@ def func(a: ti.i16) -> ti.f32:
     assert func(255) == 3.0
 
 
-@ti.test(exclude=[ti.opengl, ti.vulkan])
+@ti.test(exclude=[ti.opengl])
 def test_arg_read():
     x = ti.field(ti.i32, shape=())