[Target] Add tag for Metal - Apple M1 Max (tlc-pack#27)

vinx13 · Jan 28, 2023 · 4b0874a · 4b0874a
1 parent f1bd8a9
commit 4b0874a
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/src/target/tag.cc b/src/target/tag.cc
@@ -376,4 +376,19 @@ TVM_REGISTER_TAG_AWS_C5("aws/cpu/c5.24xlarge", 48, "cascadelake");
 
 #undef TVM_REGISTER_TAG_AWS_C5
 
+#define TVM_REGISTER_METAL_TAG(Name, ThreadsPerBlock, SharedMem, WarpSize)       \
+  TVM_REGISTER_TARGET_TAG(Name).set_config(                                      \
+      {{"kind", String("metal")},                                                \
+       {"max_threads_per_block", Integer(ThreadsPerBlock)},                      \
+       {"max_shared_memory_per_block", Integer(SharedMem)},                      \
+       {"thread_warp_size", Integer(WarpSize)},                                  \
+       {"host", Map<String, ObjectRef>{{"kind", String("llvm")},                 \
+                                       {"mtriple", String("arm64-apple-macos")}, \
+                                       {"mcpu", String("apple-latest")}}}});
+
+TVM_REGISTER_METAL_TAG("apple/m1", 1024, 32768, 32);
+TVM_REGISTER_METAL_TAG("apple/m1-max", 1024, 32768, 32);
+
+#undef TVM_REGISTER_METAL_TAG
+
 }  // namespace tvm