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