This is an automated email from the ASF dual-hosted git repository.
tqchen 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 b59da5904c [Target] Add target tags for Apple Silicon GPU (#14068)
b59da5904c is described below
commit b59da5904c8d3702b44e5f8aa09c5761bf329b7e
Author: Ruihang Lai <[email protected]>
AuthorDate: Tue Feb 21 21:08:42 2023 -0500
[Target] Add target tags for Apple Silicon GPU (#14068)
As we are recently conducting experiments on GPU of Apple Silicon, we
find it helpful to introduce the target tag for Apple Silicon GPUs, as
specialized Metal targets.
At this moment, we find the size of shared memory per thread block, the
maximum number of threads per thread block and the warp size useful,
so each specialized chip is configured by these three parameters.
The feature numbers come from Apple's official document:
https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf
---
src/target/tag.cc | 15 +++++++++++++++
1 file changed, 15 insertions(+)
diff --git a/src/target/tag.cc b/src/target/tag.cc
index c9f2414581..34a939183b 100644
--- 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_GPU_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_GPU_TAG("apple/m1-gpu", 1024, 32768, 32);
+TVM_REGISTER_METAL_GPU_TAG("apple/m2-gpu", 1024, 32768, 32);
+
+#undef TVM_REGISTER_METAL_TAG
+
} // namespace tvm
