chore: import upstream snapshot with attribution

tests/python/codegen/test_target_codegen_llvm_vla.py

@@ -0,0 +1,212 @@
+# 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.
+
+"""
+Codegen tests for VLA extensions
+"""
+
+import re
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tirx as T
+from tvm.target.codegen import llvm_version_major
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_codegen_vscale(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+    vscale = tvm.tirx.vscale()
+
+    @T.prim_func(s_tir=True)
+    def main(A: T.Buffer((5,), "int32")):
+        for i in range(5):
+            A[i] = 2 * vscale
+
+    with tvm.target.Target(target):
+        build_mod = tvm.tirx.build(main)
+
+    llvm = build_mod.inspect_source()
+    assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_buffer_load_store(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        B = T.match_buffer(b, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        B[T.ramp(0, 1, 4 * T.vscale())] = A[T.ramp(0, 1, 4 * T.vscale())]
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    assert re.findall(r"load <vscale x 4 x float>", llvm), "No scalable load in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skipif(
+    llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_scalable_broadcast(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def my_func(a: T.handle):
+        A = T.match_buffer(a, (128,), "float32")
+        T.func_attr({"global_symbol": "my_module", "tirx.noalias": True})
+        A[T.ramp(0, 1, 4 * T.vscale())] = T.broadcast(1, 4 * T.vscale())
+
+    with tvm.target.Target(target):
+        mod = tvm.tirx.build(my_func)
+
+    llvm = mod.inspect_source("ll")
+    # Older LLVM versions print the broadcast as a shufflevector of an insertelement,
+    # newer ones print it as a splat constant.
+    assert (
+        "shufflevector (<vscale x 4 x float> insertelement (<vscale x 4 x float>" in llvm
+        or "store <vscale x 4 x float> splat (float 1.000000e+00)" in llvm
+    ), "No scalable broadcast in generated LLVM."
+    assert re.findall(r" store <vscale x 4 x float>", llvm), "No scalable store in generated LLVM."
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_get_active_lane_mask(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle):
+        A = T.match_buffer(a, (30,), "int1")
+        for i in range(T.ceildiv(30, T.vscale() * 4)):
+            A[i : i + T.vscale() * 4] = T.get_active_lane_mask("uint1xvscalex4", i, 30)
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+
+
+@pytest.mark.skip(
+    reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
+)
+@pytest.mark.parametrize(
+    "target",
+    [
+        {"kind": "llvm", "mtriple": "aarch64-linux-gnu", "mattr": ["+sve"]},
+        {
+            "kind": "llvm",
+            "device": "riscv_cpu",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mattr": ["+64bit", "+a", "+c", "+d", "+f", "+m", "+v"],
+        },
+    ],
+)
+def test_predicated_scalable_buffer(target):
+    if not tvm.testing.device_enabled(target):
+        pytest.skip(f"{target} not enabled")
+
+    @T.prim_func(s_tir=True)
+    def before(a: T.handle, b: T.handle):
+        A = T.match_buffer(a, (16,), "float32")
+        B = T.match_buffer(b, (16,), "float32")
+        T.func_attr({"global_symbol": "main", "tirx.noalias": True})
+        for i_0 in T.serial(T.ceildiv(16, 4 * T.vscale())):
+            for i_1 in T.vectorized(4 * T.vscale()):
+                if i_0 * 4 * T.vscale() + i_1 < 14:
+                    B[i_0 * 4 * T.vscale() + i_1] = A[i_0 * 4 * T.vscale() + i_1] + 1.0
+
+    with tvm.target.Target(target):
+        out = tvm.tirx.build(before)
+
+    ll = out.inspect_source("ll")
+    assert "get.active.lane.mask" in ll
+    assert "llvm.masked.load" in ll
+    assert "llvm.masked.store" in ll
+
+
+if __name__ == "__main__":
+    tvm.testing.main()