[auto_scheduler] metal default hardware params (apache#7022)

electriclilies · Dec 3, 2020 · 965a67e · 965a67e
1 parent 3afde62
commit 965a67e
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/src/auto_scheduler/search_task.cc b/src/auto_scheduler/search_task.cc
@@ -71,6 +71,17 @@ HardwareParams HardwareParamsNode::GetDefaultHardwareParams(const Target& target
 
     p_hardware_params->max_vthread_extent = p_hardware_params->warp_size / 4;
 
+    return hardware_params;
+  } else if (target->kind->device_type == kDLMetal) {
+    // Reference: https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
+    // This setting looks working for Metal GPUs later than A10
+    auto hardware_params = HardwareParams(-1, 16, 64);
+    auto* p_hardware_params = hardware_params.CopyOnWrite();
+    p_hardware_params->max_shared_memory_per_block = 32 * 1024;
+    p_hardware_params->max_registers_per_block = 4 * 1024;
+    p_hardware_params->max_threads_per_block = 1024;
+    p_hardware_params->warp_size = 8;
+    p_hardware_params->max_vthread_extent = p_hardware_params->warp_size / 4;
     return hardware_params;
   } else {
     LOG(FATAL) << "No default hardware parameters for target: " << target;