(tvm) branch main updated: [Device][Metal] Fix metal warp size (#16192)

[ju...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

junrushao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 2eb17fa87f [Device][Metal] Fix metal warp size (#16192)
2eb17fa87f is described below

commit 2eb17fa87f5661cf25ebf516cd405ef7be34fc40
Author: Hongyi Jin <ji...@gmail.com>
AuthorDate: Sat Dec 2 03:29:42 2023 -0500

    [Device][Metal] Fix metal warp size (#16192)
    
    Metal warp size should be 1 on x86 device, but 32 on M1/M2 device
---
 src/runtime/metal/metal_device_api.mm | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/runtime/metal/metal_device_api.mm b/src/runtime/metal/metal_device_api.mm
index c0fe0b76a0..f7c2976d22 100644
--- a/src/runtime/metal/metal_device_api.mm
+++ b/src/runtime/metal/metal_device_api.mm
@@ -55,8 +55,14 @@ void MetalWorkspace::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) {
         break;
       }
       case kWarpSize: {
-        // Set warp size to be 1 for safty reason.
+#if defined(__x86_64__)
         *rv = 1;
+#elif defined(__aarch64__)
+        *rv = 32;
+#else
+        LOG(WARNING) << "The CPU architecture is neither x86 nor aarch64. Fallback to warp size 1.";
+        *rv = 1;
+#endif
         break;
       }
       case kMaxSharedMemoryPerBlock: