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,275 @@
+# 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.
+# pylint: disable=missing-docstring
+"""Tests for CPU DLight Reduction schedule rule."""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm import te, tirx, topi
+from tvm.s_tir import dlight as dl
+from tvm.s_tir.dlight.cpu import Reduction
+from tvm.target import Target
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _llvm_target():
+    return Target({"kind": "llvm"})
+
+
+def _rvv_target():
+    return Target(
+        {
+            "kind": "llvm",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mabi": "lp64d",
+            "mattr": ["+64bit", "+m", "+a", "+f", "+d", "+c", "+v"],
+        }
+    )
+
+
+def _build_softmax(batch, features, fast=False):
+    A = te.placeholder((batch, features), dtype="float32", name="A")
+    B = topi.nn.fast_softmax(A, axis=1) if fast else topi.nn.softmax(A, axis=1)
+    func = te.create_prim_func([A, B])
+    return tvm.IRModule({"main": func})
+
+
+def _apply_and_check(mod, target):
+    """Apply Reduction rule and verify it was applied."""
+    rule = Reduction()
+    result = rule.apply(mod["main"], target, False)
+    assert result is not None, "Reduction rule should apply"
+    return result
+
+
+# ---------------------------------------------------------------------------
+# Test: schedule applicability
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("fast", [False, True], ids=["softmax", "fast_softmax"])
+@pytest.mark.parametrize(
+    "batch,features",
+    [
+        (1, 10),
+        (1, 128),
+        (14, 185),
+        (32, 256),
+        (64, 512),
+        (128, 1024),
+        (1, 30522),
+    ],
+)
+def test_reduction_applies(batch, features, fast):
+    """Reduction rule should apply to softmax/fast_softmax of various shapes."""
+    mod = _build_softmax(batch, features, fast=fast)
+    target = _llvm_target()
+    _apply_and_check(mod, target)
+
+
+# ---------------------------------------------------------------------------
+# Test: scheduled TIR structure
+# ---------------------------------------------------------------------------
+
+
+def test_softmax_schedule_structure():
+    """Verify the scheduled TIR has expected structure:
+    - parallel on batch axis
+    - vectorized innermost loops for injective blocks
+    - split+unroll for reduction blocks
+    """
+    mod = _build_softmax(14, 185, fast=False)
+    target = _llvm_target()
+    sch = _apply_and_check(mod, target)
+    scheduled_mod = sch.mod
+
+    # Check that tirx.is_scheduled is NOT set (only set by ApplyDefaultSchedule)
+    # but the schedule should be valid
+    assert scheduled_mod is not None
+
+    # Verify via ApplyDefaultSchedule path
+    with target:
+        scheduled = dl.ApplyDefaultSchedule(Reduction())(mod)
+    func = scheduled["main"]
+
+    # Check tirx.is_scheduled is set
+    assert func.attrs and func.attrs.get("tirx.is_scheduled", False)
+
+
+def test_fast_softmax_schedule_structure():
+    """fast_softmax should keep T_fast_exp as a separate vectorizable block."""
+    mod = _build_softmax(14, 185, fast=True)
+    target = _llvm_target()
+    sch = _apply_and_check(mod, target)
+    script = str(sch.mod)
+
+    # fast_exp block should exist (not inlined)
+    assert "T_fast_exp" in script or "T_softmax_delta" in script
+    # Should have T.parallel
+    assert "T.parallel" in script
+    # Should have T.vectorized
+    assert "T.vectorized" in script
+
+
+# ---------------------------------------------------------------------------
+# Test: LLVM IR quality (cross-compile to RISC-V RVV)
+# ---------------------------------------------------------------------------
+
+
+def _codegen_llvm_ir(mod, target):
+    """Lower and codegen to LLVM IR (no linking)."""
+    bound = tirx.transform.BindTarget(target.with_host(target))(mod)
+    pipeline, finalize_host, _ = tirx.get_tir_pipeline("default")
+    lowered = pipeline(bound)
+    from tvm.tirx.build import split_host_device_mods
+
+    host_mod, _ = split_host_device_mods(lowered)
+    host_mod = finalize_host()(host_mod)
+    built = tvm.target.codegen.build_module(host_mod, target)
+    return built.inspect_source("ll")
+
+
+def _codegen_asm(mod, target):
+    """Lower and codegen to assembly (no linking)."""
+    bound = tirx.transform.BindTarget(target.with_host(target))(mod)
+    pipeline, finalize_host, _ = tirx.get_tir_pipeline("default")
+    lowered = pipeline(bound)
+    from tvm.tirx.build import split_host_device_mods
+
+    host_mod, _ = split_host_device_mods(lowered)
+    host_mod = finalize_host()(host_mod)
+    built = tvm.target.codegen.build_module(host_mod, target)
+    return built.inspect_source("s")
+
+
+@pytest.mark.parametrize("fast", [False, True], ids=["softmax", "fast_softmax"])
+def test_rvv_code_size_reduction(fast):
+    """Scheduled RVV code should be smaller than unscheduled.
+
+    The original issue (apache/tvm#18569) shows RVV softmax is 1.34x slower
+    than scalar, partly due to LLVM generating bloated code with excessive
+    unrolling. The schedule should reduce code size significantly.
+    """
+    target = _rvv_target()
+    mod = _build_softmax(14, 185, fast=fast)
+
+    # Unscheduled
+    ir_unsched = _codegen_llvm_ir(mod, target)
+    n_unsched = len(ir_unsched.splitlines())
+
+    # Scheduled
+    with target:
+        mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
+    ir_sched = _codegen_llvm_ir(mod_sched, target)
+    n_sched = len(ir_sched.splitlines())
+
+    # Scheduled should be meaningfully smaller (at least 30% reduction)
+    ratio = n_sched / n_unsched
+    assert ratio < 0.75, (
+        f"Expected >=25% code reduction, got {(1 - ratio) * 100:.1f}% "
+        f"({n_unsched} -> {n_sched} lines)"
+    )
+
+
+# The arith analyzer no longer proves vscale-bearing inequalities via
+# substitution (CanProveVscaleExpressionFromKnownValues was deleted). This
+# weakens simplification of scalable-vector index expressions, which can
+# prevent the RVV vectorization schedule from producing scalable vector ops.
+@pytest.mark.xfail(reason="arith no longer proves vscale-bearing inequalities via substitution")
+def test_rvv_fast_softmax_vectorizes_exp():
+    """fast_softmax + schedule should produce RVV vector instructions
+    for the polynomial exp approximation (no scalar exp calls)."""
+    target = _rvv_target()
+    mod = _build_softmax(14, 185, fast=True)
+    with target:
+        mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
+    ir = _codegen_llvm_ir(mod_sched, target)
+
+    # Should have zero scalar exp calls (fast_exp uses polynomial)
+    scalar_exp = sum(1 for line in ir.splitlines() if "llvm.exp.f32" in line)
+    assert scalar_exp == 0, f"Expected 0 scalar exp calls, got {scalar_exp}"
+
+    # Should have scalable vector operations
+    n_svec = ir.count("<vscale x")
+    assert n_svec > 0, "Expected scalable vector operations in LLVM IR"
+
+
+def test_rvv_asm_instruction_reduction():
+    """Scheduled RVV assembly should have fewer total instructions
+    than both unscheduled RVV and scalar RV."""
+    rvv = _rvv_target()
+    rv = Target(
+        {
+            "kind": "llvm",
+            "mtriple": "riscv64-linux-gnu",
+            "mcpu": "generic-rv64",
+            "mabi": "lp64d",
+            "mattr": ["+64bit", "+m", "+a", "+f", "+d", "+c"],
+        }
+    )
+
+    mod = _build_softmax(14, 185, fast=True)
+
+    # Scalar baseline
+    asm_rv = _codegen_asm(mod, rv)
+    n_rv = len(
+        [
+            line
+            for line in asm_rv.splitlines()
+            if line.strip() and not line.strip().startswith((".", "#", "/"))
+        ]
+    )
+
+    # RVV unscheduled
+    asm_rvv = _codegen_asm(mod, rvv)
+    n_rvv = len(
+        [
+            line
+            for line in asm_rvv.splitlines()
+            if line.strip() and not line.strip().startswith((".", "#", "/"))
+        ]
+    )
+
+    # RVV scheduled
+    with rvv:
+        mod_sched = dl.ApplyDefaultSchedule(Reduction())(mod)
+    asm_sched = _codegen_asm(mod_sched, rvv)
+    n_sched = len(
+        [
+            line
+            for line in asm_sched.splitlines()
+            if line.strip() and not line.strip().startswith((".", "#", "/"))
+        ]
+    )
+
+    # Scheduled should be smaller than both unscheduled RVV and scalar
+    assert n_sched < n_rvv, (
+        f"Scheduled ({n_sched}) should have fewer instructions than unscheduled RVV ({n_rvv})"
+    )
+    assert n_sched <= n_rv * 1.1, (
+        f"Scheduled ({n_sched}) should not be much larger than scalar RV ({n_rv})"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])