chore: import upstream snapshot with attribution

python/tvm/backend/metal/__init__.py

@@ -0,0 +1,89 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal-owned TIRx modules."""
+
+from importlib import import_module
+from pathlib import Path
+
+from tvm_ffi.libinfo import load_lib_ctypes
+
+from tvm.base import _LOADED_LIBS
+
+_LAZY_SUBMODULES = {"op", "script", "target_tags"}
+
+
+def _detect_target_from_device(dev):
+    from tvm.target import Target  # pylint: disable=import-outside-toplevel
+
+    return Target(
+        {
+            "kind": "metal",
+            "max_shared_memory_per_block": 32768,
+            "max_threads_per_block": dev.max_threads_per_block,
+            "thread_warp_size": dev.warp_size,
+        }
+    )
+
+
+def register_backend():
+    """Register Metal-owned Python semantics."""
+    from tvm.target.detect_target import register_device_target_detector
+    from tvm.tirx.script.builder import ir as builder_ir  # pylint: disable=import-outside-toplevel
+
+    runtime_dir = Path(_LOADED_LIBS["tvm_runtime"]._name).resolve().parent
+    try:
+        # Runtime sidecars only need registration side effects; libtvm_runtime is global.
+        _LOADED_LIBS["tvm_runtime_metal"] = load_lib_ctypes(
+            package="tvm",
+            target_name="tvm_runtime_metal",
+            extra_lib_paths=[runtime_dir],
+            mode="RTLD_LOCAL",
+        )
+    except (OSError, FileNotFoundError, RuntimeError):
+        pass
+    register_device_target_detector("metal", _detect_target_from_device)
+    for name, namespace in script_namespaces().items():
+        builder_ir.register_script_namespace(name, namespace)
+    import_module(f"{__name__}.target_tags")
+
+
+def script_namespaces(**_):
+    """Return Metal-owned TVMScript namespaces."""
+    from .script import MetalNamespace  # pylint: disable=import-outside-toplevel
+
+    return {"metal": MetalNamespace()}
+
+
+def script_namespace(**kwargs):
+    """Return the Metal TVMScript namespace object."""
+    return script_namespaces(**kwargs)["metal"]
+
+
+def __getattr__(name: str):
+    if name in _LAZY_SUBMODULES:
+        return import_module(f"{__name__}.{name}")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    "op",
+    "register_backend",
+    "script",
+    "script_namespace",
+    "script_namespaces",
+    "target_tags",
+]

python/tvm/backend/metal/op.py

@@ -0,0 +1,84 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal-owned TIR intrinsic builders."""
+
+from __future__ import annotations
+
+from tvm.tirx.op import call_intrin
+
+
+def make_filled_simdgroup_matrix(d, index, value, col=8, row=8):
+    """Create a filled SIMDGroup matrix."""
+
+    return call_intrin("void", "tirx.make_filled_simdgroup_matrix", d, index, value, col, row)
+
+
+def simdgroup_load(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+    """Load data from device or threadgroup memory to simdgroup."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_load",
+        d,
+        index,
+        ptr,
+        stride,
+        col,
+        row,
+        transpose_matrix,
+    )
+
+
+def simdgroup_store(d, index, ptr, stride, col=8, row=8, transpose_matrix=False):
+    """Store data from simdgroup to device or threadgroup memory."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_store",
+        d,
+        index,
+        ptr,
+        stride,
+        col,
+        row,
+        transpose_matrix,
+    )
+
+
+def simdgroup_multiply_accumulate(d, index_d, a, index_a, b, index_b, c, index_c):
+    """Multiply and accumulate two matrices in simdgroup."""
+
+    return call_intrin(
+        "void",
+        "tirx.simdgroup_multiply_accumulate",
+        d,
+        index_d,
+        a,
+        index_a,
+        b,
+        index_b,
+        c,
+        index_c,
+    )
+
+
+__all__ = [
+    "make_filled_simdgroup_matrix",
+    "simdgroup_load",
+    "simdgroup_multiply_accumulate",
+    "simdgroup_store",
+]

python/tvm/backend/metal/script.py

@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Metal TVMScript namespace."""
+
+from __future__ import annotations
+
+from tvm.backend.metal import op as _metal_op
+from tvm.tirx import Buffer
+from tvm.tirx import op as _tir_op
+from tvm.tirx.script.builder.ir import _op_wrapper
+
+
+class MetalNamespace:
+    """The Metal intrinsics submodule."""
+
+    def __init__(self):
+        self.make_filled_simdgroup_matrix = _op_wrapper(_metal_op.make_filled_simdgroup_matrix)
+        self.simdgroup_load = _op_wrapper(_metal_op.simdgroup_load)
+        self.simdgroup_store = _op_wrapper(_metal_op.simdgroup_store)
+        self.simdgroup_multiply_accumulate = _op_wrapper(_metal_op.simdgroup_multiply_accumulate)
+
+    @staticmethod
+    def simd_shuffle(var, lane):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle", var, lane)
+
+    @staticmethod
+    def simd_shuffle_up(var, delta):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_up", var, delta)
+
+    @staticmethod
+    def simd_shuffle_down(var, delta):
+        if isinstance(var, Buffer):
+            var = var[0]
+        return _tir_op.call_intrin(var.ty, "tirx.metal.simd_shuffle_down", var, delta)
+
+
+__all__ = ["MetalNamespace"]

python/tvm/backend/metal/target_tags.py

@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Apple Metal GPU target tags."""
+
+from tvm.target import register_tag
+
+_METAL_HOST_TRIPLE = "arm64-apple-macos"
+
+
+def _register_metal_tag(name, max_threads, shared_mem, warp_size, mcpu):
+    try:
+        from tvm.target.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,
+        {
+            "kind": "metal",
+            "max_threads_per_block": max_threads,
+            "max_shared_memory_per_block": shared_mem,
+            "thread_warp_size": warp_size,
+            "host": {
+                "kind": "llvm",
+                "mtriple": _METAL_HOST_TRIPLE,
+                "mcpu": mcpu,
+            },
+        },
+    )
+
+
+_register_metal_tag("apple/m1-gpu", 1024, 32768, 32, "apple-m1")
+_register_metal_tag("apple/m1-gpu-restricted", 256, 32768, 32, "apple-m1")
+_register_metal_tag("apple/m2-gpu", 1024, 32768, 32, "apple-m2")