From 40bae7523a345e27599081147a8322e7e4e2d6c3 Mon Sep 17 00:00:00 2001 From: Ruihang Lai Date: Tue, 21 Feb 2023 13:53:31 -0500 Subject: [PATCH] [Target] Add target tags for Apple Silicon GPU 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 c9f24145814b..34a939183be4 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{{"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