chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,151 @@
+# 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.
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm import relax
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+
+def test_pipeline_compile():
+    target = tvm.target.Target("llvm", host="llvm")
+    pipeline = relax.pipeline.get_default_pipeline(target)
+
+    @tvm.script.ir_module
+    class Mod:
+        @R.function
+        def main(x: R.Tensor((3, 4), "float32"), y: R.Tensor((3, 4), "float32")):
+            lv0 = R.add(x, y)
+            return lv0
+
+    mod = Mod
+    mod = pipeline(mod)
+
+    ex = tvm.compile(mod, target)
+    x_np = np.random.rand(3, 4).astype(np.float32)
+    y_np = np.random.rand(3, 4).astype(np.float32)
+    x = tvm.runtime.tensor(x_np)
+    y = tvm.runtime.tensor(y_np)
+
+    vm = relax.VirtualMachine(ex, tvm.cpu())
+    z = vm["main"](x, y)
+    tvm.testing.assert_allclose(z.numpy(), x_np + y_np, rtol=1e-7, atol=1e-7)
+
+
+def test_pipeline_with_kv_cache():
+    """A dummy pipline that simulates KV update."""
+    target = tvm.target.Target("llvm", host="llvm")
+    pipeline = relax.pipeline.get_default_pipeline(target)
+
+    @tvm.script.ir_module
+    class Mod:
+        @R.function
+        def create_kv_cache(reserve_slots: R.Shape(["m"])):
+            # just allocate minimum slot since it is only used to signal dtype
+            m = T.int64()
+            init_data = R.ones((1, 4), "float32")
+            kv_cache = R.call_pure_packed(
+                "vm.builtin.attention_kv_cache_create",
+                init_data,
+                R.shape([m, 4]),
+                0,
+                ty_args=[R.Any()],
+            )
+            return kv_cache
+
+        @R.function(pure=False)
+        def main(
+            x: R.Tensor((1, 4), "float32"),
+            y: R.Tensor((1, 4), "float32"),
+            shape: R.Shape(["L", 4]),
+            kv_cache: R.Any,
+        ):
+            L = T.int64()
+            # computation of the current value
+            curr_value = R.add(x, y)
+            # update cache
+            kv_cache = R.call_packed(
+                "vm.builtin.attention_kv_cache_append", kv_cache, curr_value, ty_args=[R.Any]
+            )
+            # return the updated cache view
+            kv = R.call_packed(
+                "vm.builtin.attention_kv_cache_view",
+                kv_cache,
+                shape,
+                ty_args=[R.Tensor((L, 4), "float32")],
+            )
+            return (kv, kv_cache)
+
+    mod = Mod
+    mod = pipeline(mod)
+
+    ex = tvm.compile(mod, target)
+
+    num_steps = 8
+    cache_np = np.empty((num_steps, 4), dtype="float32")
+    vm = relax.VirtualMachine(ex, tvm.cpu())
+
+    kv_cache = vm["create_kv_cache"](tvm_ffi.Shape([1]))
+
+    for i in range(num_steps):
+        x_np = np.random.rand(1, 4).astype(np.float32)
+        y_np = np.random.rand(1, 4).astype(np.float32)
+        x = tvm.runtime.tensor(x_np)
+        y = tvm.runtime.tensor(y_np)
+        np_shape = (i + 1, 4)
+        kv, kv_cache = vm["main"](x, y, tvm_ffi.Shape(np_shape), kv_cache)
+
+        cache_np[i, :] = x_np + y_np
+        tvm.testing.assert_allclose(kv.numpy(), cache_np[: np_shape[0], :], rtol=1e-7, atol=1e-7)
+
+
+@pytest.mark.parametrize("target_name", ["vulkan", "webgpu"])
+@pytest.mark.parametrize(
+    "pipeline_func",
+    [
+        relax.pipeline.library_dispatch_passes,
+        relax.pipeline.legalize_passes,
+        relax.pipeline.dataflow_lower_passes,
+        relax.pipeline.finalize_passes,
+        relax.pipeline.get_default_pipeline,
+    ],
+)
+def test_gpu_generic_fallback(target_name, pipeline_func):
+    target = tvm.target.Target(target_name)
+    result = pipeline_func(target)
+    assert result is not None
+
+
+@pytest.mark.parametrize("target_name", ["hexagon", "c"])
+@pytest.mark.parametrize(
+    "pipeline_func",
+    [
+        relax.pipeline.library_dispatch_passes,
+        relax.pipeline.legalize_passes,
+        relax.pipeline.dataflow_lower_passes,
+        relax.pipeline.finalize_passes,
+        relax.pipeline.get_default_pipeline,
+    ],
+)
+def test_non_gpu_target_raises_error(target_name, pipeline_func):
+    target = tvm.target.Target(target_name)
+    with pytest.raises(ValueError, match="not yet supported"):
+        pipeline_func(target)