Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions python/tvm/target/codegen.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,24 @@ def llvm_get_vector_width(target=None):
return _ffi_api.llvm_get_vector_width(target)


def llvm_is_valid_cpu(cpu, triple):
"""Check if a CPU name is valid for the given LLVM triple.

Parameters
----------
cpu : str
The CPU name to check (e.g. "apple-m1").
triple : str
The LLVM target triple (e.g. "arm64-apple-macos").

Returns
-------
is_valid : bool
True if the CPU name is recognized by LLVM for the given triple.
"""
return _ffi_api.llvm_is_valid_cpu(cpu, triple)


def llvm_version_major(allow_none=False):
"""Get the major LLVM version.

Expand Down
11 changes: 10 additions & 1 deletion python/tvm/target/tag_registry/metal.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,17 @@

from .registry import register_tag

_METAL_HOST_TRIPLE = "arm64-apple-macos"


def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu):
try:
from ..codegen import llvm_is_valid_cpu

if not llvm_is_valid_cpu(mcpu, _METAL_HOST_TRIPLE):
return
except Exception: # pylint: disable=broad-except
pass # LLVM not available; register unconditionally
register_tag(
name,
{
Expand All @@ -29,7 +38,7 @@ def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu):
"thread_warp_size": warp_size,
"host": {
"kind": "llvm",
"mtriple": "arm64-apple-macos",
"mtriple": _METAL_HOST_TRIPLE,
"mcpu": mcpu,
},
},
Expand Down
9 changes: 9 additions & 0 deletions src/target/llvm/llvm_module.cc
Original file line number Diff line number Diff line change
Expand Up @@ -743,6 +743,15 @@ static void LLVMReflectionRegister() {
LLVMTargetInfo llvm_target(*llvm_instance, use_target);
return llvm_target.TargetHasCPUFeature(feature);
})
.def("target.llvm_is_valid_cpu",
[](ffi::String cpu, ffi::String triple) -> bool {
auto llvm_instance = std::make_unique<LLVMInstance>();
ffi::Map<ffi::String, ffi::Any> target_map;
target_map.Set("kind", ffi::String("llvm"));
target_map.Set("mtriple", triple);
LLVMTargetInfo llvm_backend(*llvm_instance, Target(target_map));
return llvm_backend.IsValidCPU(std::string(cpu));
})
.def("target.llvm_version_major", []() -> int { return TVM_LLVM_VERSION / 10; })
.def("ffi.Module.load_from_file.ll",
[](std::string filename, std::string fmt) -> ffi::Module {
Expand Down
Loading