Skip to content

Commit

Permalink
[Target] Add target tags for Apple Silicon GPU
Browse files Browse the repository at this point in the history
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
  • Loading branch information
MasterJH5574 committed Feb 21, 2023
1 parent 347d79c commit 40bae75
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions src/target/tag.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit 40bae75

Please sign in to comment.