From 3763f0a192feea1d67fd4a0e3b0fab8d90fe3dc2 Mon Sep 17 00:00:00 2001 From: Anders Lie Date: Sat, 12 Sep 2026 14:31:36 -0700 Subject: [PATCH] [Metal] Resolve target attributes from device properties Query each Metal device once for the facts that decide what codegen may emit: the highest Metal Shading Language version it compiles, bfloat16, SIMD-group permute/reduction/matrix support, and the Metal 4 family. Expose them through "device_api.metal.get_target_property", mirroring the Vulkan device API, and register the matching attributes on the "metal" target kind so Target.from_device("metal") reports them. Serve the threadgroup memory limit through GetAttr instead of hard-coding 32 KB in target detection. Carry the MSL version on the Metal module so runtime compilation uses the version codegen generated for, instead of choosing one from the device at compile time. The module byte format gains one field after "fmt". --- python/tvm/target/detect_target.py | 11 +- src/runtime/metal/metal_common.h | 38 +++++ src/runtime/metal/metal_device_api.mm | 133 +++++++++++++++++- src/runtime/metal/metal_module.mm | 58 ++++---- src/target/metal/codegen_metal.cc | 7 +- src/target/metal/metal_fallback_module.cc | 19 ++- src/target/metal/metal_fallback_module.h | 15 +- src/target/target_kind.cc | 8 ++ .../codegen/test_target_codegen_metal.py | 64 +++++++++ 9 files changed, 307 insertions(+), 46 deletions(-) diff --git a/python/tvm/target/detect_target.py b/python/tvm/target/detect_target.py index 81accfed1287..28f655aaff4c 100644 --- a/python/tvm/target/detect_target.py +++ b/python/tvm/target/detect_target.py @@ -23,12 +23,21 @@ def _detect_metal(dev: Device) -> Target: + f_get_target_property = get_global_func("device_api.metal.get_target_property") return Target( { "kind": "metal", - "max_shared_memory_per_block": 32768, + "max_shared_memory_per_block": dev.max_shared_memory_per_block, "max_threads_per_block": dev.max_threads_per_block, "thread_warp_size": dev.warp_size, + "metal_language_version": f_get_target_property(dev, "metal_language_version"), + "supports_bfloat16": f_get_target_property(dev, "supports_bfloat16"), + "supports_simdgroup_permute": f_get_target_property(dev, "supports_simdgroup_permute"), + "supports_simdgroup_reduction": f_get_target_property( + dev, "supports_simdgroup_reduction" + ), + "supports_simdgroup_matrix": f_get_target_property(dev, "supports_simdgroup_matrix"), + "supports_metal4": f_get_target_property(dev, "supports_metal4"), } ) diff --git a/src/runtime/metal/metal_common.h b/src/runtime/metal/metal_common.h index 101eb3a2d585..eb5c65593801 100644 --- a/src/runtime/metal/metal_common.h +++ b/src/runtime/metal/metal_common.h @@ -305,6 +305,35 @@ class MetalRawStream final : public Stream { }; +/*! + * \brief Device facts that decide which "metal" target attributes a device supports. + * + * Queried once per device when the workspace initializes. Python target + * detection reads them through "device_api.metal.get_target_property" to + * populate the matching attributes of the "metal" target kind, the same way + * the Vulkan device API reports its VulkanDeviceProperties. + */ +struct MetalDeviceProperties { + /*! + * \brief Highest Metal Shading Language version the device can compile, + * encoded as major * 10 + minor (for example 31 for MSL 3.1). + */ + int metal_language_version; + bool supports_bfloat16; + bool supports_simdgroup_permute; + bool supports_simdgroup_reduction; + bool supports_simdgroup_matrix; + bool supports_metal4; +}; + +/*! + * \brief Map a metal_language_version attribute value to the MTLLanguageVersion enum. + * \param version The version encoded as major * 10 + minor. + * \return The matching enum value. + * \note Throws when this build's SDK does not know the requested version. + */ +MTLLanguageVersion MetalLanguageVersionFromNumber(int version); + /*! * \brief Process global Metal workspace. */ @@ -314,6 +343,8 @@ class MetalWorkspace final : public DeviceAPI { std::vector> devices; // Warp size constant std::vector warp_size; + // Target-relevant properties of each device, parallel to `devices`. + std::vector device_properties; MetalWorkspace(); // Destructor ~MetalWorkspace(); @@ -327,6 +358,13 @@ class MetalWorkspace final : public DeviceAPI { // override device API void SetDevice(Device dev) final; void GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) final; + /*! + * \brief Report one "metal" target attribute supported by the device. + * \param dev The device to query. + * \param property The target attribute name, e.g. "supports_bfloat16". + * \param rv The value; left unset when the property is unknown. + */ + void GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv); void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final; void FreeDataSpace(Device dev, void* ptr) final; TVMStreamHandle CreateStream(Device dev) final; diff --git a/src/runtime/metal/metal_device_api.mm b/src/runtime/metal/metal_device_api.mm index 2a44d98f109f..1b07577f2364 100644 --- a/src/runtime/metal/metal_device_api.mm +++ b/src/runtime/metal/metal_device_api.mm @@ -42,6 +42,100 @@ return inst; } +namespace { + +// Highest Metal Shading Language version the running OS can compile, encoded +// as major * 10 + minor. The SDK guards keep older toolchains building; the +// @available checks keep the answer correct on older systems. +int MaximumMetalLanguageVersion() { + int version = 23; +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 120000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 150000) + if (@available(macOS 12.0, iOS 15.0, *)) version = 24; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 130000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 160000) + if (@available(macOS 13.0, iOS 16.0, *)) version = 30; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 140000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 170000) + if (@available(macOS 14.0, iOS 17.0, *)) version = 31; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 150000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 180000) + if (@available(macOS 15.0, iOS 18.0, *)) version = 32; +#endif + return version; +} + +bool SupportsFamily(id device, MTLGPUFamily family) { + if (@available(macOS 10.15, iOS 13.0, *)) return [device supportsFamily:family]; + return false; +} + +MetalDeviceProperties QueryDeviceProperties(id device) { + // Feature availability follows the Metal feature set tables: + // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf + const bool apple6 = SupportsFamily(device, MTLGPUFamilyApple6); + const bool apple7 = SupportsFamily(device, MTLGPUFamilyApple7); + const bool mac2 = SupportsFamily(device, MTLGPUFamilyMac2); + bool metal4 = false; +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 260000) + if (@available(macOS 26.0, iOS 26.0, *)) { + metal4 = [device supportsFamily:MTLGPUFamilyMetal4]; + } +#endif + // MSL 4.0 is only offered to devices in the Metal 4 family; every other + // device stops at the highest 3.x version the OS provides. + const int language_version = metal4 ? 40 : MaximumMetalLanguageVersion(); + MetalDeviceProperties properties; + properties.metal_language_version = language_version; + properties.supports_bfloat16 = language_version >= 31 && (apple6 || mac2); + properties.supports_simdgroup_permute = apple6 || mac2; + properties.supports_simdgroup_reduction = apple7 || mac2; + properties.supports_simdgroup_matrix = apple7; + properties.supports_metal4 = metal4; + return properties; +} + +} // namespace + +MTLLanguageVersion MetalLanguageVersionFromNumber(int version) { + switch (version) { + case 23: + return MTLLanguageVersion2_3; +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 120000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 150000) + case 24: + return MTLLanguageVersion2_4; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 130000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 160000) + case 30: + return MTLLanguageVersion3_0; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 140000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 170000) + case 31: + return MTLLanguageVersion3_1; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 150000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 180000) + case 32: + return MTLLanguageVersion3_2; +#endif +#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000) || \ + (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 260000) + case 40: + return MTLLanguageVersion4_0; +#endif + default: + TVM_FFI_THROW(RuntimeError) << "Metal language version " << version + << " is not available in this build"; + } +} + void MetalWorkspace::GetAttr(Device dev, DeviceAttrKind kind, ffi::Any* rv) { AUTORELEASEPOOL { size_t index = static_cast(dev.device_id); @@ -66,8 +160,10 @@ #endif break; } - case kMaxSharedMemoryPerBlock: - return; + case kMaxSharedMemoryPerBlock: { + *rv = static_cast([devices[dev.device_id] maxThreadgroupMemoryLength]); + break; + } case kComputeVersion: return; case kDeviceName: @@ -102,6 +198,31 @@ }; } +void MetalWorkspace::GetTargetProperty(Device dev, const std::string& property, ffi::Any* rv) { + size_t index = static_cast(dev.device_id); + TVM_FFI_ICHECK_LT(index, device_properties.size()) << "Invalid device id " << index; + const MetalDeviceProperties& properties = device_properties[index]; + + if (property == "metal_language_version") { + *rv = int64_t(properties.metal_language_version); + } + if (property == "supports_bfloat16") { + *rv = properties.supports_bfloat16; + } + if (property == "supports_simdgroup_permute") { + *rv = properties.supports_simdgroup_permute; + } + if (property == "supports_simdgroup_reduction") { + *rv = properties.supports_simdgroup_reduction; + } + if (property == "supports_simdgroup_matrix") { + *rv = properties.supports_simdgroup_matrix; + } + if (property == "supports_metal4") { + *rv = properties.supports_metal4; + } +} + static const char* kDummyKernel = R"A0B0( using namespace metal; // Simple copy kernel @@ -163,6 +284,7 @@ int GetWarpSize(id dev) { // on iPhone id d = MTLCreateSystemDefaultDevice(); devices.push_back(d); + device_properties.push_back(QueryDeviceProperties(d)); #else NSArray >* devs = MTLCopyAllDevices(); for (size_t i = 0; i < devs.count; ++i) { @@ -170,6 +292,7 @@ int GetWarpSize(id dev) { devices.push_back(d); DLOG(INFO) << "Intializing Metal device " << i << ", name=" << [d.name UTF8String]; warp_size.push_back(GetWarpSize(d)); + device_properties.push_back(QueryDeviceProperties(d)); } #endif this->ReinitializeDefaultStreams(); @@ -400,6 +523,12 @@ int GetWarpSize(id dev) { DeviceAPI* ptr = MetalWorkspace::Global(); *rv = static_cast(ptr); }) + .def("device_api.metal.get_target_property", + [](Device dev, const std::string& property) { + ffi::Any rv; + MetalWorkspace::Global()->GetTargetProperty(dev, property, &rv); + return rv; + }) .def("metal.ResetGlobalState", []() { MetalWorkspace::Global()->ReinitializeDefaultStreams(); }) .def("metal.GetProfileCounters", diff --git a/src/runtime/metal/metal_module.mm b/src/runtime/metal/metal_module.mm index 7a51998c5320..73775fed420d 100644 --- a/src/runtime/metal/metal_module.mm +++ b/src/runtime/metal/metal_module.mm @@ -45,26 +45,12 @@ #include "metal_common.h" #include "tvm/runtime/device_api.h" -#if (defined(__MAC_OS_X_VERSION_MAX_ALLOWED) && __MAC_OS_X_VERSION_MAX_ALLOWED >= 260000) || \ - (defined(__IPHONE_OS_VERSION_MAX_ALLOWED) && __IPHONE_OS_VERSION_MAX_ALLOWED >= 260000) -#define TVM_METAL_HAS_MSL_4_0 1 -#endif - namespace tvm { namespace runtime { /*! \brief Maximum number of GPU supported in MetalModule. */ static constexpr const int kMetalMaxNumDevice = 32; -static bool MetalDeviceSupportsMetal4(id device) { -#if defined(TVM_METAL_HAS_MSL_4_0) - if (@available(macOS 26.0, iOS 26.0, *)) { - return [device supportsFamily:MTLGPUFamilyMetal4]; - } -#endif - return false; -} - // Module to support thread-safe multi-GPU execution. // The runtime will contain a per-device module table // The modules will be lazily loaded @@ -74,14 +60,17 @@ static bool MetalDeviceSupportsMetal4(id device) { // src/target/metal/metal_fallback_module.h. The per-kernel `smap` // payload is Map regardless of whether the format is // text MSL ("metal") or compiled metallib ("metallib") — text vs binary - // distinction lives in `fmt`. + // distinction lives in `fmt`. `metal_language_version` is the MSL version + // codegen assumed (major * 10 + minor); fmt="metal" sources are compiled + // with exactly that version on every device. MetalModuleNode(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source) + ffi::Map source, int metal_language_version) : smap_(std::move(smap)), fmt_(std::move(fmt)), fmap_(std::move(fmap)), - source_(std::move(source)) {} + source_(std::move(source)), + metal_language_version_(metal_language_version) {} const char* kind() const final { return "metal"; } @@ -93,14 +82,16 @@ int GetPropertyMask() const final { ffi::Optional GetFunction(const ffi::String& name) final; ffi::Bytes SaveToBytes() const final { - // 3 fields [fmt][fmap][smap]. Source map is in-memory inspection only - // and is NEVER serialized — matches the cross-backend rule. + // 4 fields [fmt][metal_language_version][fmap][smap]. Source map is + // in-memory inspection only and is NEVER serialized — matches the + // cross-backend rule. // MetalFallbackModuleNode::SaveToBytes (in // src/target/metal/metal_fallback_module.cc) MUST mirror this format // byte-for-byte; see one-way comment there. std::string result; support::BytesOutStream stream(&result); stream.Write(fmt_); + stream.Write(metal_language_version_); stream.Write(fmap_); stream.Write(smap_); return ffi::Bytes(std::move(result)); @@ -138,14 +129,13 @@ int GetPropertyMask() const final { const ffi::Bytes& source = (*kernel).second; if (fmt_ == "metal") { + const int device_language_version = + w->device_properties[device_id].metal_language_version; + TVM_FFI_ICHECK_LE(metal_language_version_, device_language_version) + << "Metal module was generated for MSL " << metal_language_version_ << " but device " + << device_id << " compiles at most MSL " << device_language_version; MTLCompileOptions* opts = [[MTLCompileOptions alloc] init]; - MTLLanguageVersion language_version = MTLLanguageVersion2_3; -#if defined(TVM_METAL_HAS_MSL_4_0) - if (MetalDeviceSupportsMetal4(w->devices[device_id])) { - language_version = MTLLanguageVersion4_0; - } -#endif - opts.languageVersion = language_version; + opts.languageVersion = metal::MetalLanguageVersionFromNumber(metal_language_version_); opts.fastMathEnabled = YES; // Per-kernel payload is bytes; treat as UTF-8 MSL source. std::string source_str(source.data(), source.size()); @@ -211,6 +201,8 @@ int GetPropertyMask() const final { ffi::Map fmap_; // In-memory source map for InspectSource — never serialized. ffi::Map source_; + // MSL version the kernels were generated for (major * 10 + minor). + int metal_language_version_; // function information. std::vector finfo_; // internal mutex when updating the module @@ -330,11 +322,12 @@ void operator()(ffi::PackedArgs args, ffi::Any* rv, const ArgUnion64* pack_args) static ffi::Module MetalModuleCreateImpl(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source) { + ffi::Map source, + int metal_language_version) { ffi::ObjectPtr n; AUTORELEASEPOOL { n = ffi::make_object(std::move(smap), std::move(fmt), std::move(fmap), - std::move(source)); + std::move(source), metal_language_version); }; return ffi::Module(n); } @@ -342,14 +335,16 @@ void operator()(ffi::PackedArgs args, ffi::Any* rv, const ArgUnion64* pack_args) static ffi::Module MetalModuleLoadFromBytes(const ffi::Bytes& bytes) { support::BytesInStream stream(bytes); ffi::String fmt; + int metal_language_version; ffi::Map fmap; ffi::Map smap; stream.Read(&fmt); + TVM_FFI_ICHECK(stream.Read(&metal_language_version)); TVM_FFI_ICHECK(stream.Read(&fmap)); stream.Read(&smap); // Source map is not serialized — reconstructed empty on load. return MetalModuleCreateImpl(std::move(smap), std::move(fmt), std::move(fmap), - ffi::Map()); + ffi::Map(), metal_language_version); } void SetMetalStream(TVMStreamHandle stream) { @@ -371,9 +366,10 @@ void SetMetalStream(TVMStreamHandle stream) { .def("ffi.Module.load_from_bytes.metal", MetalModuleLoadFromBytes) .def("ffi.Module.create.metal", [](ffi::Map smap, ffi::String fmt, - ffi::Map fmap, ffi::Map source) { + ffi::Map fmap, ffi::Map source, + int metal_language_version) { return MetalModuleCreateImpl(std::move(smap), std::move(fmt), std::move(fmap), - std::move(source)); + std::move(source), metal_language_version); }) .def("metal.SetStream", SetMetalStream); } diff --git a/src/target/metal/codegen_metal.cc b/src/target/metal/codegen_metal.cc index c84df824a14f..3c918274d11e 100644 --- a/src/target/metal/codegen_metal.cc +++ b/src/target/metal/codegen_metal.cc @@ -486,8 +486,13 @@ ffi::Module BuildMetal(IRModule mod, Target target) { // map keyed by "metal" — only used by InspectSource and never serialized. ffi::Map source; source.Set("metal", source_maker.str()); + // The module records the MSL version this target generated so the runtime + // compiles the source with the same language version on every device. + const int metal_language_version = + target->GetAttr("metal_language_version").value_or(Integer(23))->value; return target::MetalModuleCreateWithFallback(std::move(smap), ffi::String(fmt), - ExtractFuncInfo(mod), std::move(source)); + ExtractFuncInfo(mod), std::move(source), + metal_language_version); } TVM_FFI_STATIC_INIT_BLOCK() { diff --git a/src/target/metal/metal_fallback_module.cc b/src/target/metal/metal_fallback_module.cc index 0954c307b4cb..cc377bae4c99 100644 --- a/src/target/metal/metal_fallback_module.cc +++ b/src/target/metal/metal_fallback_module.cc @@ -42,11 +42,12 @@ class MetalFallbackModuleNode : public ffi::ModuleObj { public: MetalFallbackModuleNode(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source) + ffi::Map source, int metal_language_version) : smap_(std::move(smap)), fmt_(std::move(fmt)), fmap_(std::move(fmap)), - source_(std::move(source)) {} + source_(std::move(source)), + metal_language_version_(metal_language_version) {} // Mirror the real module's kind so consumers cannot distinguish at the // kind/api layer. Saved bytes load back as a real MetalModuleNode on a @@ -72,11 +73,13 @@ class MetalFallbackModuleNode : public ffi::ModuleObj { // mirror the change here. The dependency is one-way: this file // follows; metal_module.mm does not reference this file. // - // 3 fields only — the source map is in-memory inspection material and - // is NEVER serialized (matches upstream behavior for all backends). + // 4 fields [fmt][metal_language_version][fmap][smap] — the source map is + // in-memory inspection material and is NEVER serialized (matches + // upstream behavior for all backends). std::string buffer; support::BytesOutStream stream(&buffer); stream.Write(fmt_); + stream.Write(metal_language_version_); stream.Write(fmap_); stream.Write(smap_); return ffi::Bytes(std::move(buffer)); @@ -117,13 +120,17 @@ class MetalFallbackModuleNode : public ffi::ModuleObj { ffi::Map fmap_; // In-memory source map for InspectSource — never serialized. ffi::Map source_; + // MSL version the kernels were generated for (major * 10 + minor). + int metal_language_version_; }; ffi::Module MetalFallbackModuleCreate(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source) { + ffi::Map source, + int metal_language_version) { auto n = ffi::make_object(std::move(smap), std::move(fmt), - std::move(fmap), std::move(source)); + std::move(fmap), std::move(source), + metal_language_version); return ffi::Module(n); } diff --git a/src/target/metal/metal_fallback_module.h b/src/target/metal/metal_fallback_module.h index b30ebd991289..e5031023a19b 100644 --- a/src/target/metal/metal_fallback_module.h +++ b/src/target/metal/metal_fallback_module.h @@ -51,7 +51,8 @@ namespace target { */ ffi::Module MetalFallbackModuleCreate(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source); + ffi::Map source, + int metal_language_version); /*! * \brief Codegen-time Metal module factory. Tries the FFI-registered @@ -65,23 +66,27 @@ ffi::Module MetalFallbackModuleCreate(ffi::Map smap, ff * - `TVM_COMPILE_FORCE_FALLBACK` env var truthy → fallback regardless of * registry state (used by per-backend fallback tests on a USE_X=ON CI * box). + * + * `metal_language_version` is the MSL version codegen generated for + * (major * 10 + minor); the runtime compiles fmt="metal" sources with it. */ inline ffi::Module MetalModuleCreateWithFallback(ffi::Map smap, ffi::String fmt, ffi::Map fmap, - ffi::Map source) { + ffi::Map source, + int metal_language_version) { if (tvm::support::GetEnv("TVM_COMPILE_FORCE_FALLBACK", false)) { return MetalFallbackModuleCreate(std::move(smap), std::move(fmt), std::move(fmap), - std::move(source)); + std::move(source), metal_language_version); } // Registry: "ffi.Module.create.metal" — real Metal runtime factory. // Grep hint: grep -rn 'ffi.Module.create.metal' src/ auto fcreate = ffi::Function::GetGlobal("ffi.Module.create.metal"); if (fcreate.has_value()) { - return (*fcreate)(smap, fmt, fmap, source).cast(); + return (*fcreate)(smap, fmt, fmap, source, metal_language_version).cast(); } return MetalFallbackModuleCreate(std::move(smap), std::move(fmt), std::move(fmap), - std::move(source)); + std::move(source), metal_language_version); } } // namespace target diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc index f7af7bdbcb49..083bce088cb7 100644 --- a/src/target/target_kind.cc +++ b/src/target/target_kind.cc @@ -431,6 +431,14 @@ TVM_REGISTER_TARGET_KIND("metal", kDLMetal) .add_attr_option("max_shared_memory_per_block", refl::DefaultValue(32768)) .add_attr_option("thread_warp_size", refl::DefaultValue(16)) .add_attr_option("max_function_args", refl::DefaultValue(31)) + // Metal Shading Language version to generate and compile, as major * 10 + minor. + .add_attr_option("metal_language_version", refl::DefaultValue(23)) + // Feature support, populated from the device by Target.from_device + .add_attr_option("supports_bfloat16") + .add_attr_option("supports_simdgroup_permute") + .add_attr_option("supports_simdgroup_reduction") + .add_attr_option("supports_simdgroup_matrix") + .add_attr_option("supports_metal4") .set_default_keys({"metal", "gpu"}); TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan) diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 4f8ab4efdd87..6f5b52bb64bd 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. import numpy as np +import pytest import tvm import tvm.testing @@ -234,5 +235,68 @@ def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")): host_lib.export_library(lib_path) +@tvm.testing.requires_gpu +@tvm.testing.requires_metal +def test_target_from_device_reports_device_properties(): + dev = tvm.metal() + target = tvm.target.Target.from_device(dev) + + assert int(target.attrs["max_shared_memory_per_block"]) == dev.max_shared_memory_per_block + assert int(target.attrs["max_threads_per_block"]) == dev.max_threads_per_block + assert int(target.attrs["thread_warp_size"]) == dev.warp_size + assert int(target.attrs["metal_language_version"]) >= 23 + for name in ( + "supports_bfloat16", + "supports_simdgroup_permute", + "supports_simdgroup_reduction", + "supports_simdgroup_matrix", + "supports_metal4", + ): + assert name in target.attrs, name + if bool(target.attrs["supports_bfloat16"]): + assert int(target.attrs["metal_language_version"]) >= 31 + if bool(target.attrs["supports_metal4"]): + assert int(target.attrs["metal_language_version"]) == 40 + + +@tvm.testing.requires_gpu +@tvm.testing.requires_metal +def test_module_compiles_with_target_language_version(tmp_path): + n = 16 + + @I.ir_module + class Module: + @T.prim_func + def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")): + T.func_attr({"tirx.noalias": True}) + for i in T.thread_binding(n, thread="threadIdx.x"): + with T.sblock("B"): + v_i = T.axis.spatial(n, i) + T.reads(A[v_i]) + T.writes(B[v_i]) + B[v_i] = A[v_i] + 1.0 + + dev = tvm.metal() + a = np.arange(n).astype("float32") + + # The version travels with the module through export and load. + target = tvm.target.Target({"kind": "metal", "metal_language_version": 30}) + lib_path = str(tmp_path / "lib.so") + tvm.compile(Module, target=target).export_library(lib_path) + loaded = tvm.runtime.load_module(lib_path) + a_nd = tvm.runtime.tensor(a, dev) + b_nd = tvm.runtime.empty((n,), "float32", dev) + loaded["main"](a_nd, b_nd) + tvm.testing.assert_allclose(b_nd.numpy(), a + 1.0, atol=1e-5, rtol=1e-5) + + # A module generated for a newer MSL than the device compiles is rejected. + detected = tvm.target.Target.from_device(dev) + too_new = int(detected.attrs["metal_language_version"]) + 1 + target = tvm.target.Target({"kind": "metal", "metal_language_version": too_new}) + f = tvm.compile(Module, target=target) + with pytest.raises(Exception, match="compiles at most MSL"): + f(a_nd, b_nd) + + if __name__ == "__main__": tvm.testing.main()