diff --git a/.github/workflows/ci_sim.yml b/.github/workflows/ci_sim.yml index 498a730505..d6807f14d9 100644 --- a/.github/workflows/ci_sim.yml +++ b/.github/workflows/ci_sim.yml @@ -60,6 +60,7 @@ jobs: - 'test/vpto/**' - 'test/tilelang_st/**' - 'test/dsl-st/**' + - 'test/tilelib-st/**' - 'tilelang-dsl/**' - 'ptodsl/**' - 'tools/ptoas/**' @@ -125,37 +126,33 @@ jobs: fetch-depth: 1 persist-credentials: false - - name: Resolve LLVM directories + - name: Resolve simulator environment shell: bash run: | set -euo pipefail - echo "LLVM_ROOT=${RUNNER_TOOL_CACHE}/llvm-project" >> "${GITHUB_ENV}" - echo "LLVM_DIR=${RUNNER_TOOL_CACHE}/llvm-project/llvm/build-assert" >> "${GITHUB_ENV}" - echo "MLIR_PYTHONPATH=${RUNNER_TOOL_CACHE}/llvm-project/llvm/build-assert/tools/mlir/python_packages/mlir_core" >> "${GITHUB_ENV}" - - name: Resolve LLVM cache key - id: llvm-cache-key - shell: bash - run: | - set -euo pipefail - # Resolve to a Git object that ls-remote can handle: either a tag - # (LLVM_TAG) or a branch head (LLVM_REF). Only one is expected. - ref="${LLVM_TAG:-${LLVM_REF}}" - sha="$(git ls-remote "${LLVM_REPO}" "${ref}" | awk '{print $1}')" - if [[ -z "${sha}" ]]; then - echo "ERROR: failed to resolve ${LLVM_REPO} ${ref}" >&2 + detect_ascend_home() { + for d in \ + "${ASCEND_HOME_PATH:-}" \ + /usr/local/Ascend/cann \ + /usr/local/Ascend/cann-* \ + /usr/local/Ascend/ascend-toolkit/latest + do + [[ -n "${d}" && -d "${d}" ]] || continue + printf '%s\n' "${d}" + return 0 + done + return 1 + } + + ASCEND_HOME_PATH_DETECTED="$(detect_ascend_home || true)" + if [[ -z "${ASCEND_HOME_PATH_DETECTED}" ]]; then + echo "ERROR: failed to detect ASCEND_HOME_PATH on self-hosted runner" >&2 exit 1 fi - echo "sha=${sha}" >> "${GITHUB_OUTPUT}" - echo "key=llvm-build-${sha}-assert-v1" >> "${GITHUB_OUTPUT}" - - name: Restore LLVM build cache - id: llvm-cache - continue-on-error: true - uses: actions/cache/restore@v4 - with: - path: ${{ env.LLVM_DIR }} - key: ${{ steps.llvm-cache-key.outputs.key }} + echo "ASCEND_HOME_PATH=${ASCEND_HOME_PATH_DETECTED}" >> "${GITHUB_ENV}" + echo "PTOAS_BIN=${GITHUB_WORKSPACE}/build/tools/ptoas/ptoas" >> "${GITHUB_ENV}" - name: Ensure runner dependencies shell: bash @@ -201,12 +198,54 @@ jobs: python3 -m pip install setuptools wheel 'pybind11<3' nanobind numpy ml-dtypes fi + - name: Resolve LLVM directories + shell: bash + run: | + set -euo pipefail + echo "LLVM_ROOT=${RUNNER_TOOL_CACHE}/llvm-project" >> "${GITHUB_ENV}" + echo "LLVM_DIR=${RUNNER_TOOL_CACHE}/llvm-project/llvm/build-assert" >> "${GITHUB_ENV}" + echo "MLIR_PYTHONPATH=${RUNNER_TOOL_CACHE}/llvm-project/llvm/build-assert/tools/mlir/python_packages/mlir_core" >> "${GITHUB_ENV}" + + - name: Resolve PTODSL build directories + shell: bash + run: | + set -euo pipefail + echo "PTO_DSL_ST_BUILD_DIR=${GITHUB_WORKSPACE}/build-ptodsl" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_INSTALL_DIR=${GITHUB_WORKSPACE}/install-ptodsl" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_PTOAS_BIN=${GITHUB_WORKSPACE}/build-ptodsl/tools/ptoas/ptoas" >> "${GITHUB_ENV}" + + - name: Resolve LLVM cache key + id: llvm-cache-key + shell: bash + run: | + set -euo pipefail + # Resolve to a Git object that ls-remote can handle: either a tag + # (LLVM_TAG) or a branch head (LLVM_REF). Only one is expected. + ref="${LLVM_TAG:-${LLVM_REF}}" + sha="$(git ls-remote "${LLVM_REPO}" "${ref}" | awk '{print $1}')" + if [[ -z "${sha}" ]]; then + echo "ERROR: failed to resolve ${LLVM_REPO} ${ref}" >&2 + exit 1 + fi + echo "sha=${sha}" >> "${GITHUB_OUTPUT}" + echo "key=llvm-build-${sha}-assert-v1" >> "${GITHUB_OUTPUT}" + + - name: Restore LLVM build cache + id: llvm-cache + continue-on-error: true + uses: actions/cache/restore@v4 + with: + path: ${{ env.LLVM_DIR }} + key: ${{ steps.llvm-cache-key.outputs.key }} + - name: Clean CI work dirs shell: bash run: | set -euo pipefail rm -rf "${GITHUB_WORKSPACE}/build" + rm -rf "${GITHUB_WORKSPACE}/build-ptodsl" rm -rf "${PTO_INSTALL_DIR}" + rm -rf "${PTO_DSL_ST_INSTALL_DIR}" rm -rf "${VPTO_SIM_WORKSPACE}" rm -rf "${TILELANG_DSL_WORKSPACE}" rm -rf "${PYPTO_WORKSPACE}" @@ -245,8 +284,8 @@ jobs: # Clean the build directory so that stale generated files from a # previous run (e.g. _smt_ops_gen.py left behind when the ref # changed) do not leak into the fresh build. - rm -rf llvm/build-assert - cmake -G Ninja -S llvm -B llvm/build-assert \ + rm -rf "${LLVM_DIR}" + cmake -G Ninja -S llvm -B "${LLVM_DIR}" \ -DLLVM_ENABLE_PROJECTS="mlir;clang" \ -DBUILD_SHARED_LIBS=ON \ -DLLVM_ENABLE_ASSERTIONS=ON \ @@ -258,7 +297,7 @@ jobs: -DCMAKE_BUILD_TYPE=Release \ -DLLVM_TARGETS_TO_BUILD="host" - ninja -C llvm/build-assert + ninja -C "${LLVM_DIR}" - name: Save LLVM build cache if: steps.llvm-cache.outputs.cache-hit != 'true' @@ -281,34 +320,6 @@ jobs: PTO_INSTALL_DIR="${PTO_INSTALL_DIR}" \ python3 -m pip install . --no-build-isolation --no-deps --ignore-installed --prefix "${PTO_INSTALL_DIR}" - - name: Resolve simulator environment - shell: bash - run: | - set -euo pipefail - - detect_ascend_home() { - for d in \ - "${ASCEND_HOME_PATH:-}" \ - /usr/local/Ascend/cann \ - /usr/local/Ascend/cann-* \ - /usr/local/Ascend/ascend-toolkit/latest - do - [[ -n "${d}" && -d "${d}" ]] || continue - printf '%s\n' "${d}" - return 0 - done - return 1 - } - - ASCEND_HOME_PATH_DETECTED="$(detect_ascend_home || true)" - if [[ -z "${ASCEND_HOME_PATH_DETECTED}" ]]; then - echo "ERROR: failed to detect ASCEND_HOME_PATH on self-hosted runner" >&2 - exit 1 - fi - - echo "ASCEND_HOME_PATH=${ASCEND_HOME_PATH_DETECTED}" >> "${GITHUB_ENV}" - echo "PTOAS_BIN=${GITHUB_WORKSPACE}/build/tools/ptoas/ptoas" >> "${GITHUB_ENV}" - - name: Checkout PyPTO uses: actions/checkout@v4 with: @@ -505,18 +516,169 @@ jobs: echo "SIM_LIB_DIR=${SIM_LIB_DIR}" >> "${GITHUB_ENV}" echo "SIM_LIB_DIR=${SIM_LIB_DIR}" - - name: Run VPTO SIM validation - if: ${{ true }} + - name: Resolve PTODSL Python shell: bash run: | set -euo pipefail - mkdir -p "${VPTO_SIM_WORKSPACE}" - WORK_SPACE="${VPTO_SIM_WORKSPACE}" \ - ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ - PTOAS_BIN="${PTOAS_BIN}" \ - DEVICE=SIM \ - JOBS="${JOBS:-32}" \ - bash test/vpto/scripts/run_host_vpto_validation_parallel.sh + source "${ASCEND_HOME_PATH}/bin/setenv.bash" + + add_python_candidate() { + local candidate="$1" + local resolved + [[ -n "${candidate}" ]] || return 0 + if [[ "${candidate}" != */* ]]; then + resolved="$(command -v "${candidate}" 2>/dev/null || true)" + else + resolved="${candidate}" + fi + [[ -n "${resolved}" && -x "${resolved}" ]] || return 0 + resolved="$(readlink -f "${resolved}")" + case ":${PYTHON_CANDIDATES[*]}:" in + *":${resolved}:"*) return 0 ;; + esac + PYTHON_CANDIDATES+=("${resolved}") + } + + has_torch_npu_packages() { + "$1" - <<'PY' + import importlib.util + + missing = [ + name for name in ("torch", "torch_npu") + if importlib.util.find_spec(name) is None + ] + raise SystemExit(1 if missing else 0) + PY + } + + probe_ptodsl_runtime_python() { + TORCH_DEVICE_BACKEND_AUTOLOAD=0 "$1" - <<'PY' + import sys + import nanobind + import numpy + import pybind11 + import yaml + import torch + import torch_npu # noqa: F401 + + print(sys.executable) + print(f"python {sys.version_info.major}.{sys.version_info.minor}") + print("torch", torch.__version__) + print("torch_npu", getattr(torch_npu, "__version__", "unknown")) + print("numpy", numpy.__version__) + print("nanobind", getattr(nanobind, "__version__", "unknown")) + print("pybind11", pybind11.__version__) + PY + } + + missing_ptodsl_python_deps() { + "$1" - <<'PY' + import importlib.util + + deps = [ + ("setuptools", "setuptools"), + ("wheel", "wheel"), + ("nanobind", "nanobind"), + ("numpy", "numpy"), + ("ml_dtypes", "ml-dtypes"), + ("yaml", "PyYAML"), + ] + + missing = [ + requirement for module_name, requirement in deps + if importlib.util.find_spec(module_name) is None + ] + + try: + import pybind11 + version = tuple(int(part) for part in pybind11.__version__.split(".")[:1]) + if version >= (3,): + missing.append("pybind11<3") + except Exception: + missing.append("pybind11<3") + + print(" ".join(missing)) + PY + } + + PYTHON_CANDIDATES=() + if [[ -n "${PTO_DSL_ST_PYTHON_BIN:-}" ]]; then + add_python_candidate "${PTO_DSL_ST_PYTHON_BIN}" + else + add_python_candidate python3 + shopt -s nullglob + for candidate in \ + /home/*/miniconda3/bin/python3 \ + /home/*/miniconda3/bin/python \ + /home/*/anaconda3/bin/python3 \ + /home/*/anaconda3/bin/python \ + /home/*/miniconda3/envs/*/bin/python \ + /home/*/anaconda3/envs/*/bin/python \ + /opt/conda/envs/*/bin/python + do + add_python_candidate "${candidate}" + done + shopt -u nullglob + fi + + SELECTED_BASE_PYTHON="" + for candidate in "${PYTHON_CANDIDATES[@]}"; do + echo "Probing PTODSL base Python package presence: ${candidate}" + if has_torch_npu_packages "${candidate}"; then + SELECTED_BASE_PYTHON="${candidate}" + break + fi + done + + if [[ -z "${SELECTED_BASE_PYTHON}" ]]; then + echo "ERROR: PTODSL DSL ST requires an existing Python runtime with torch and torch_npu." >&2 + echo "ERROR: this workflow intentionally does not install torch or torch_npu on every run." >&2 + echo "ERROR: set PTO_DSL_ST_PYTHON_BIN to a compatible pre-installed interpreter." >&2 + exit 1 + fi + + PYTHON_ABI_TAG="$("${SELECTED_BASE_PYTHON}" - <<'PY' + import sys + print(f"py{sys.version_info.major}{sys.version_info.minor}") + PY + )" + BASE_HASH="$(printf '%s' "${SELECTED_BASE_PYTHON}" | sha256sum | cut -c1-12)" + PYTHON_TAG="${PYTHON_ABI_TAG}-${BASE_HASH}" + PTODSL_PYTHON_ROOT="${RUNNER_TOOL_CACHE}/ptodsl-python/${PYTHON_TAG}" + if [[ ! -x "${PTODSL_PYTHON_ROOT}/bin/python" ]]; then + rm -rf "${PTODSL_PYTHON_ROOT}" + "${SELECTED_BASE_PYTHON}" -m venv --system-site-packages "${PTODSL_PYTHON_ROOT}" + fi + + PTODSL_PYTHON="${PTODSL_PYTHON_ROOT}/bin/python" + if ! "${PTODSL_PYTHON}" -m pip --version >/dev/null 2>&1; then + "${PTODSL_PYTHON}" -m ensurepip --upgrade + fi + missing_deps="$(missing_ptodsl_python_deps "${PTODSL_PYTHON}")" + if [[ -n "${missing_deps}" ]]; then + "${PTODSL_PYTHON}" -m pip install ${missing_deps} + fi + + probe_ptodsl_runtime_python "${PTODSL_PYTHON}" + + PTO_DSL_ST_LLVM_DIR="${RUNNER_TOOL_CACHE}/llvm-project/llvm/build-assert-${PYTHON_TAG}" + PTO_DSL_ST_MLIR_PYTHONPATH="${PTO_DSL_ST_LLVM_DIR}/tools/mlir/python_packages/mlir_core" + PTO_DSL_ST_PYTHON_SITE="$( + PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" "${PTODSL_PYTHON}" - <<'PY' + import os + import sysconfig + + prefix = os.environ["PTO_INSTALL_DIR"] + print(sysconfig.get_path("purelib", vars={"base": prefix, "platbase": prefix})) + PY + )" + + echo "PTO_DSL_ST_BASE_PYTHON=${SELECTED_BASE_PYTHON}" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_PYTHON_BIN=${PTODSL_PYTHON}" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_PYTHON_TAG=${PYTHON_TAG}" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_LLVM_DIR=${PTO_DSL_ST_LLVM_DIR}" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_MLIR_PYTHONPATH=${PTO_DSL_ST_MLIR_PYTHONPATH}" >> "${GITHUB_ENV}" + echo "PTO_DSL_ST_PYTHON_SITE=${PTO_DSL_ST_PYTHON_SITE}" >> "${GITHUB_ENV}" - name: Run TileLang DSL CI shell: bash @@ -537,15 +699,190 @@ jobs: 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/run_ci.log" fi - - name: Run PTODSL DSL ST CI + - name: Restore PTODSL LLVM build cache + id: ptodsl-llvm-cache + continue-on-error: true + uses: actions/cache/restore@v4 + with: + path: ${{ env.PTO_DSL_ST_LLVM_DIR }} + key: llvm-build-${{ steps.llvm-cache-key.outputs.sha }}-assert-${{ env.PTO_DSL_ST_PYTHON_TAG }}-v2 + + - name: Build PTODSL LLVM/MLIR + if: steps.ptodsl-llvm-cache.outputs.cache-hit != 'true' + shell: bash + run: | + set -euo pipefail + cd "${LLVM_ROOT}" + export CC=gcc + export CXX=g++ + PTODSL_PYTHON_PREFIX="$(dirname "$(dirname "${PTO_DSL_ST_PYTHON_BIN}")")" + export VIRTUAL_ENV="${PTODSL_PYTHON_PREFIX}" + export PATH="${PTODSL_PYTHON_PREFIX}/bin:${PATH}" + PYBIND11_CMAKE_DIR="$("${PTO_DSL_ST_PYTHON_BIN}" -m pybind11 --cmakedir)" + NANOBIND_CMAKE_DIR="$("${PTO_DSL_ST_PYTHON_BIN}" -m nanobind --cmake_dir)" + PYTHON_INCLUDE_DIR="$("${PTO_DSL_ST_PYTHON_BIN}" - <<'PY' + import sysconfig + print(sysconfig.get_path("include")) + PY + )" + PYTHON_LIBRARY="$("${PTO_DSL_ST_PYTHON_BIN}" - <<'PY' + import sysconfig + print(sysconfig.get_config_var("LIBDIR") + "/" + sysconfig.get_config_var("LDLIBRARY")) + PY + )" + rm -rf "${PTO_DSL_ST_LLVM_DIR}" + cmake -G Ninja -S llvm -B "${PTO_DSL_ST_LLVM_DIR}" \ + -DLLVM_ENABLE_PROJECTS="mlir;clang" \ + -DBUILD_SHARED_LIBS=ON \ + -DLLVM_ENABLE_ASSERTIONS=ON \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DPython3_EXECUTABLE="${PTO_DSL_ST_PYTHON_BIN}" \ + -DPython_EXECUTABLE="${PTO_DSL_ST_PYTHON_BIN}" \ + -DPython3_ROOT_DIR="${PTODSL_PYTHON_PREFIX}" \ + -DPython_ROOT_DIR="${PTODSL_PYTHON_PREFIX}" \ + -DPython3_INCLUDE_DIR="${PYTHON_INCLUDE_DIR}" \ + -DPython_INCLUDE_DIR="${PYTHON_INCLUDE_DIR}" \ + -DPython3_LIBRARY="${PYTHON_LIBRARY}" \ + -DPython_LIBRARY="${PYTHON_LIBRARY}" \ + -DPython3_FIND_STRATEGY=LOCATION \ + -DPython_FIND_STRATEGY=LOCATION \ + -DPython3_FIND_VIRTUALENV=ONLY \ + -DPython_FIND_VIRTUALENV=ONLY \ + -DPython3_FIND_REGISTRY=NEVER \ + -DPython_FIND_REGISTRY=NEVER \ + -DPython3_FIND_FRAMEWORK=NEVER \ + -DPython_FIND_FRAMEWORK=NEVER \ + -Dpybind11_DIR="${PYBIND11_CMAKE_DIR}" \ + -Dnanobind_DIR="${NANOBIND_CMAKE_DIR}" \ + -DCMAKE_PREFIX_PATH="${PYBIND11_CMAKE_DIR};${NANOBIND_CMAKE_DIR};${PTODSL_PYTHON_PREFIX}" \ + -DCMAKE_BUILD_TYPE=Release \ + -DLLVM_TARGETS_TO_BUILD="host" + + ninja -C "${PTO_DSL_ST_LLVM_DIR}" + + - name: Save PTODSL LLVM build cache + if: steps.ptodsl-llvm-cache.outputs.cache-hit != 'true' + continue-on-error: true + uses: actions/cache/save@v4 + with: + path: ${{ env.PTO_DSL_ST_LLVM_DIR }} + key: llvm-build-${{ steps.llvm-cache-key.outputs.sha }}-assert-${{ env.PTO_DSL_ST_PYTHON_TAG }}-v2 + + - name: Build PTODSL PTOAS + shell: bash + run: | + set -euo pipefail + rm -rf "${PTO_DSL_ST_BUILD_DIR}" "${PTO_DSL_ST_INSTALL_DIR}" + export CC=gcc + export CXX=g++ + LLVM_BUILD_DIR="${PTO_DSL_ST_LLVM_DIR}" \ + PTO_BUILD_DIR="${PTO_DSL_ST_BUILD_DIR}" \ + PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" \ + "${PTO_DSL_ST_PYTHON_BIN}" -m pip install . --no-build-isolation --no-deps --ignore-installed --prefix "${PTO_DSL_ST_INSTALL_DIR}" + + - name: Probe PTODSL runtime Python shell: bash run: | set -euo pipefail mkdir -p "${TILELANG_DSL_WORKSPACE}" - export LLVM_BUILD_DIR="${LLVM_DIR}" - export PYTHON_BIN="python3" + export LLVM_BUILD_DIR="${PTO_DSL_ST_LLVM_DIR}" + export MLIR_PYTHON_ROOT="${PTO_DSL_ST_MLIR_PYTHONPATH}" + export PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" + export PTO_PYTHON_BUILD_ROOT="${PTO_DSL_ST_BUILD_DIR}/python" + export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTOAS_PYTHON_SITE="${PTO_DSL_ST_PYTHON_SITE}" + export TORCH_DEVICE_BACKEND_AUTOLOAD=0 + source "${ASCEND_HOME_PATH}/bin/setenv.bash" + + probe_ptodsl_python() { + PYTHONPATH="${PTO_DSL_ST_INSTALL_DIR}:${PTOAS_PYTHON_SITE}:${PTO_DSL_ST_MLIR_PYTHONPATH}:${PTO_DSL_ST_BUILD_DIR}/python:${PYTHONPATH:-}" \ + "${PTO_DSL_ST_PYTHON_BIN}" - <<'PY' + import sys + import numpy + import torch + import torch_npu # noqa: F401 + from ptodsl import pto # noqa: F401 + from mlir.dialects import pto as _pto # noqa: F401 + + print(sys.executable) + print("torch", torch.__version__) + print("torch_npu", getattr(torch_npu, "__version__", "unknown")) + print("numpy", numpy.__version__) + PY + } + + PROBE_LOG="${TILELANG_DSL_WORKSPACE}/ptodsl-python-probe.log" + : > "${PROBE_LOG}" + echo "Probing PTODSL DSL ST Python: ${PTO_DSL_ST_PYTHON_BIN}" | tee -a "${PROBE_LOG}" + if ! probe_ptodsl_python >> "${PROBE_LOG}" 2>&1; then + cat "${PROBE_LOG}" >&2 + echo "ERROR: selected PTODSL Python cannot import torch, torch_npu, ptodsl, and MLIR PTO bindings." >&2 + exit 1 + fi + + cat "${PROBE_LOG}" + + - name: Run VPTO SIM validation + if: ${{ true }} + shell: bash + run: | + set -euo pipefail + mkdir -p "${VPTO_SIM_WORKSPACE}" + export LLVM_BUILD_DIR="${PTO_DSL_ST_LLVM_DIR}" + export MLIR_PYTHON_ROOT="${PTO_DSL_ST_MLIR_PYTHONPATH}" + export PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" + export PTO_PYTHON_BUILD_ROOT="${PTO_DSL_ST_BUILD_DIR}/python" + export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTOAS_PYTHON_SITE="${PTO_DSL_ST_PYTHON_SITE}" + export TORCH_DEVICE_BACKEND_AUTOLOAD=0 + source "${ASCEND_HOME_PATH}/bin/setenv.bash" + + WORK_SPACE="${VPTO_SIM_WORKSPACE}" \ + ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ + PTOAS_BIN="${PTO_DSL_ST_PTOAS_BIN}" \ + DEVICE=SIM \ + JOBS="${JOBS:-32}" \ + bash test/vpto/scripts/run_host_vpto_validation_parallel.sh + + - name: Run TileLib ST CI + shell: bash + run: | + set -euo pipefail + mkdir -p "${TILELANG_DSL_WORKSPACE}" + export LLVM_BUILD_DIR="${PTO_DSL_ST_LLVM_DIR}" + export MLIR_PYTHON_ROOT="${PTO_DSL_ST_MLIR_PYTHONPATH}" + export PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" + export PTO_PYTHON_BUILD_ROOT="${PTO_DSL_ST_BUILD_DIR}/python" + export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTOAS_PYTHON_SITE="${PTO_DSL_ST_PYTHON_SITE}" + export TORCH_DEVICE_BACKEND_AUTOLOAD=0 + source "${ASCEND_HOME_PATH}/bin/setenv.bash" + + ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ + PTOAS_BIN="${PTO_DSL_ST_PTOAS_BIN}" \ + scripts/sim_dsl.sh test/tilelib-st/run_tilelib_st.py -- test/tilelib-st/a5 \ + 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/tilelib-st-a5.log" + + - name: Run PTODSL ST CI + shell: bash + run: | + set -euo pipefail + mkdir -p "${TILELANG_DSL_WORKSPACE}" + export LLVM_BUILD_DIR="${PTO_DSL_ST_LLVM_DIR}" + export MLIR_PYTHON_ROOT="${PTO_DSL_ST_MLIR_PYTHONPATH}" + export PTO_INSTALL_DIR="${PTO_DSL_ST_INSTALL_DIR}" + export PTO_PYTHON_BUILD_ROOT="${PTO_DSL_ST_BUILD_DIR}/python" + export PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTO_PYTHON_BIN="${PTO_DSL_ST_PYTHON_BIN}" + export PTOAS_PYTHON_SITE="${PTO_DSL_ST_PYTHON_SITE}" + export TORCH_DEVICE_BACKEND_AUTOLOAD=0 + source "${ASCEND_HOME_PATH}/bin/setenv.bash" + ASCEND_HOME_PATH="${ASCEND_HOME_PATH}" \ - PTOAS_BIN="${PTOAS_BIN}" \ + PTOAS_BIN="${PTO_DSL_ST_PTOAS_BIN}" \ scripts/sim_dsl.sh test/dsl-st \ 2>&1 | tee "${TILELANG_DSL_WORKSPACE}/ptodsl-dsl-st.log" @@ -557,6 +894,8 @@ jobs: path: | ${{ env.TILELANG_DSL_WORKSPACE }}/run_ci.log ${{ env.TILELANG_DSL_WORKSPACE }}/ptodsl-dsl-st.log + ${{ env.TILELANG_DSL_WORKSPACE }}/tilelib-st-a5.log + ${{ env.TILELANG_DSL_WORKSPACE }}/ptodsl-python-probe.log if-no-files-found: warn - name: Run TileLang DSL unit tests diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 4520d249a8..8d043bc285 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -79,7 +79,7 @@ static bool shouldFoldAddrFamily(FoldIntrinsicMode mode) { return mode == FoldIntrinsicMode::All || mode == FoldIntrinsicMode::AddrOnly; } -static void eraseDeadAllocTileOps(func::FuncOp func) { +static bool eraseDeadAllocTileOps(func::FuncOp func) { SmallVector deadAllocs; func.walk([&](pto::AllocTileOp alloc) { if (alloc.getResult().use_empty()) @@ -88,6 +88,65 @@ static void eraseDeadAllocTileOps(func::FuncOp func) { for (pto::AllocTileOp alloc : llvm::reverse(deadAllocs)) alloc.erase(); + return !deadAllocs.empty(); +} + +static bool isPTOViewBridgeType(Type type) { + return isa(type); +} + +static bool eraseDeadViewBridgeCasts(func::FuncOp func) { + SmallVector deadCasts; + func.walk([&](UnrealizedConversionCastOp castOp) { + if (!castOp.use_empty() || castOp.getNumOperands() != 1 || + castOp.getNumResults() != 1) + return; + + Type srcTy = castOp.getOperand(0).getType(); + Type dstTy = castOp.getResult(0).getType(); + if ((isa(srcTy) && isPTOViewBridgeType(dstTy)) || + (isPTOViewBridgeType(srcTy) && isa(dstTy))) + deadCasts.push_back(castOp); + }); + + for (auto castOp : llvm::reverse(deadCasts)) + castOp.erase(); + return !deadCasts.empty(); +} + +static bool eraseDeadMemrefViewOps(func::FuncOp func) { + SmallVector deadMemrefOps; + func.walk([&](Operation *op) { + if ((isa(op) || + isa(op)) && + op->use_empty()) + deadMemrefOps.push_back(op); + }); + + for (Operation *op : llvm::reverse(deadMemrefOps)) + op->erase(); + return !deadMemrefOps.empty(); +} + +static bool eraseDeadTensorViewOps(func::FuncOp func) { + SmallVector deadViewOps; + func.walk([&](Operation *op) { + if ((isa(op) || + isa(op)) && + op->use_empty()) + deadViewOps.push_back(op); + }); + + for (Operation *op : llvm::reverse(deadViewOps)) + op->erase(); + return !deadViewOps.empty(); +} + +static void eraseDeadViewChains(func::FuncOp func) { + while (eraseDeadViewBridgeCasts(func) || eraseDeadMemrefViewOps(func) || + eraseDeadTensorViewOps(func) || eraseDeadAllocTileOps(func)) { + } } struct TileHandleInfo { @@ -728,34 +787,6 @@ struct FoldTileBufIntrinsicsPass } } - // Clean up dead unrealized_conversion_cast ops that bridged - // memref -> partition_tensor_view / tile_buf and are now unused - // after folding. - SmallVector deadCasts; - func.walk([&](UnrealizedConversionCastOp castOp) { - if (castOp.use_empty() && castOp.getNumOperands() == 1 && - isa(castOp.getOperand(0).getType()) && - isa( - castOp.getResult(0).getType())) - deadCasts.push_back(castOp); - }); - for (auto castOp : llvm::reverse(deadCasts)) - castOp.erase(); - - while (true) { - SmallVector deadMemrefOps; - func.walk([&](Operation *op) { - if ((isa(op) || - isa(op)) && - op->use_empty()) - deadMemrefOps.push_back(op); - }); - if (deadMemrefOps.empty()) - break; - for (auto *op : llvm::reverse(deadMemrefOps)) - op->erase(); - } - // Erase pto.set_validshape ops. Every valid-shape reader // (get_validshape / tile_valid_{rows,cols} / tile_buf_addr) has been // folded above, so the runtime metadata writes have no remaining @@ -788,7 +819,7 @@ struct FoldTileBufIntrinsicsPass } } - eraseDeadAllocTileOps(func); + eraseDeadViewChains(func); } }; diff --git a/ptodsl/docs/user_guide/01-introduction.md b/ptodsl/docs/user_guide/01-introduction.md index 12eba5ce58..7435052cd2 100644 --- a/ptodsl/docs/user_guide/01-introduction.md +++ b/ptodsl/docs/user_guide/01-introduction.md @@ -208,8 +208,10 @@ def my_kernel( barriers by hand, and work with raw pointers — useful when you need to hand-tune instruction schedules or overlap DMA with compute. -`mode` only affects what you can write inside the function body. It doesn't -change how you compile or launch the kernel. +For native launch builds, `mode` also selects the default PTOAS build policy: +`mode="auto"` keeps the PTOAS default build level and enables sync insertion, +while `mode="explicit"` uses `--pto-level=level3` and leaves synchronization +under user control by default. #### `backend`: VPTO vs EmitC diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md index 61e4b8e50a..a7a566624d 100644 --- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md +++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md @@ -50,8 +50,10 @@ The **`backend`** parameter selects the compilation target: rejected at decoration time with an actionable diagnostic. The **`mode`** parameter selects the programming model within the kernel body -(see Section 3.4). `mode` only affects what you can write inside the function — -it doesn't change how you compile or launch the kernel. +(see Section 3.4). For native launch builds, `mode="auto"` keeps PTOAS default +build level and enables sync insertion by default, while `mode="explicit"` uses +`--pto-level=level3` and disables sync insertion by default. This matches the +manual-address, user-managed staging contract of explicit kernels. `@pto.jit` owns compilation (tracing + lowering), caching, and — for `entry=True` — runtime launch binding. The compute-unit decorators diff --git a/ptodsl/ptodsl/_diagnostics.py b/ptodsl/ptodsl/_diagnostics.py index 572b1d5950..92226e000d 100644 --- a/ptodsl/ptodsl/_diagnostics.py +++ b/ptodsl/ptodsl/_diagnostics.py @@ -341,6 +341,15 @@ def illegal_inline_subkernel_placement_error(role: str, outer_role: str | None) ) +def subkernel_kernel_kind_mismatch_error(role: str, kernel_kind: str) -> RuntimeError: + """Return one diagnostic for mixing explicit @pto.jit kernel kind with the opposite subkernel kind.""" + return RuntimeError( + f"@pto.{role} cannot be lowered inside an explicit @pto.jit(kernel_kind={kernel_kind!r}) " + "module. Remove the explicit kernel_kind so PTOAS can split cube/vector sections, " + "or keep subkernel scopes in the same physical kind." + ) + + def inline_subkernel_value_escape_error(role: str, type_text: str) -> RuntimeError: """Return one diagnostic for outlined inline-scope values escaping their helper boundary.""" return RuntimeError( @@ -503,6 +512,7 @@ def unsupported_public_surface_error(name: str) -> AttributeError: "subkernel_host_tensor_boundary_error", "subkernel_illegal_annotation_error", "subkernel_illegal_parameter_kind_error", + "subkernel_kernel_kind_mismatch_error", "subkernel_missing_annotation_error", "subkernel_signature_boundary_error", "tile_row_alignment_error", diff --git a/ptodsl/ptodsl/_jit.py b/ptodsl/ptodsl/_jit.py index fc86a21b9f..26fb8311e5 100644 --- a/ptodsl/ptodsl/_jit.py +++ b/ptodsl/ptodsl/_jit.py @@ -34,6 +34,15 @@ _MODULE_ATTRS = ("pto.target_arch",) _SUPPORTED_FRONTEND_OPTION_KEYS = {"ast_rewrite", "rewrite_part", "dump_rewritten_source"} _SUPPORTED_REWRITE_PARTS = {"control_flow"} +_DEFAULT_KERNEL_KIND = "vector" + + +class _DefaultKernelKindSentinel: + def __repr__(self) -> str: + return repr(_DEFAULT_KERNEL_KIND) + + +_DEFAULT_KERNEL_KIND_SENTINEL = _DefaultKernelKindSentinel() def _normalize_mode(mode: str, *, fn=None) -> str: @@ -164,7 +173,7 @@ def jit( name=None, *, target: str = "a5", - kernel_kind: str = "vector", + kernel_kind: str = _DEFAULT_KERNEL_KIND_SENTINEL, backend: str = "vpto", entry: bool = True, mode: str = "auto", @@ -180,10 +189,10 @@ def jit( ---------- name: IR function name (defaults to the Python function name). target: Target architecture string, e.g. ``"a5"``. - kernel_kind: authored default physical kind, used for native build selection - and VPTO authoring intent. PTODSL now expresses physical regions - through ``pto.section.vector/cube`` instead of child-module - ``pto.kernel_kind`` attributes. + kernel_kind: optional authored physical kind, used for native build selection + and explicit single-kind VPTO authoring intent. When omitted, + PTODSL keeps the historical vector default while allowing + subkernel sections to express mixed cube/vector regions. backend: ``"vpto"`` or ``"emitc"`` – records the intended backend. entry: ``True`` for launchable kernel entries, ``False`` for helpers. mode: ``"auto"`` or ``"explicit"`` – feeds child compile policy. @@ -238,12 +247,15 @@ def decorator(fn): source_file = inspect.getsourcefile(fn) or inspect.getfile(fn) except (OSError, TypeError): source_file = None + kernel_kind_explicit = kernel_kind is not _DEFAULT_KERNEL_KIND_SENTINEL + effective_kernel_kind = kernel_kind if kernel_kind_explicit else _DEFAULT_KERNEL_KIND compiler = KernelCompiler( fn.__name__, KernelModuleSpec( function_name=fn_name, target_arch=target, - kernel_kind=kernel_kind, + kernel_kind=effective_kernel_kind, + kernel_kind_explicit=kernel_kind_explicit, backend=normalized_backend, entry=entry, mode=normalized_mode, @@ -307,6 +319,10 @@ def __ptodsl_cache_signature__(self): self._compiler._kernel_identity, module_spec.function_name, module_spec.entry, + module_spec.backend, + module_spec.mode, + module_spec.kernel_kind, + module_spec.kernel_kind_explicit, ) def _build_default_module(self): diff --git a/ptodsl/ptodsl/_runtime/cache.py b/ptodsl/ptodsl/_runtime/cache.py index 6231c16b02..9552399401 100644 --- a/ptodsl/ptodsl/_runtime/cache.py +++ b/ptodsl/ptodsl/_runtime/cache.py @@ -67,6 +67,7 @@ def write_manifest( launch_symbol: str, mlir_digest: str, launch_cpp_digest: str, + compile_config_digest: str, link_config_digest: str, ) -> None: artifacts.cache_dir.mkdir(parents=True, exist_ok=True) @@ -76,6 +77,7 @@ def write_manifest( "shared_library": str(artifacts.shared_library), "mlir_digest": mlir_digest, "launch_cpp_digest": launch_cpp_digest, + "compile_config_digest": compile_config_digest, "link_config_digest": link_config_digest, } artifacts.manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") @@ -90,6 +92,7 @@ def is_native_build_current( *, mlir_text: str, launch_cpp_text: str, + compile_config_text: str, link_config_text: str, ) -> bool: required = ( @@ -110,6 +113,7 @@ def is_native_build_current( return ( manifest.get("mlir_digest") == _content_digest(mlir_text) and manifest.get("launch_cpp_digest") == _content_digest(launch_cpp_text) + and manifest.get("compile_config_digest") == _content_digest(compile_config_text) and manifest.get("link_config_digest") == _content_digest(link_config_text) ) diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 0821c326be..410a3997ae 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -76,13 +76,34 @@ def _effective_insert_sync(*, mode: str, insert_sync: bool | None) -> bool: return mode != "explicit" +def _effective_pto_level(*, mode: str) -> str | None: + return "level3" if mode == "explicit" else None + + def _source_ptoas_overrides(module_spec) -> dict: if getattr(module_spec, "jit_source", None) is None: return {} - overrides = {"backend": module_spec.backend} - if module_spec.mode == "explicit": - overrides["pto_level"] = "level3" - return overrides + return {"backend": module_spec.backend} + + +def _compile_config_text( + *, + module_spec, + effective_insert_sync: bool, + effective_pto_level: str | None, + ptoas_overrides: dict, +) -> str: + return "\n".join( + [ + f"target_arch={module_spec.target_arch}", + f"kernel_kind={module_spec.kernel_kind}", + f"mode={module_spec.mode}", + f"insert_sync={effective_insert_sync}", + f"pto_level={effective_pto_level}", + f"backend={ptoas_overrides.get('backend')}", + "enable_tile_op_expand=True", + ] + ) def _host_compile_flags() -> list[str]: @@ -191,6 +212,18 @@ def build_native_library( ir_function_name=ir_function_name, kernel_signature=kernel_signature, ) + effective_insert_sync = _effective_insert_sync( + mode=module_spec.mode, + insert_sync=module_spec.insert_sync, + ) + effective_pto_level = _effective_pto_level(mode=module_spec.mode) + ptoas_overrides = _source_ptoas_overrides(module_spec) + compile_config_text = _compile_config_text( + module_spec=module_spec, + effective_insert_sync=effective_insert_sync, + effective_pto_level=effective_pto_level, + ptoas_overrides=ptoas_overrides, + ) sim_mode = bool(os.environ.get("MSPROF_SIMULATOR_MODE")) link_config_text = "\n".join(runtime_library_flags(sim_mode=sim_mode)) @@ -198,6 +231,7 @@ def build_native_library( artifacts, mlir_text=mlir_text, launch_cpp_text=launch_cpp_text, + compile_config_text=compile_config_text, link_config_text=link_config_text, ): return artifacts.shared_library, launch_symbol @@ -210,11 +244,9 @@ def build_native_library( artifacts.mlir_path, artifacts.kernel_object, target_arch=module_spec.target_arch, - insert_sync=_effective_insert_sync( - mode=module_spec.mode, - insert_sync=module_spec.insert_sync, - ), - **_source_ptoas_overrides(module_spec), + insert_sync=effective_insert_sync, + pto_level=effective_pto_level, + **ptoas_overrides, ) launch_object = artifacts.cache_dir / "launch.o" @@ -237,6 +269,7 @@ def build_native_library( launch_symbol=launch_symbol, mlir_digest=_content_digest(mlir_text), launch_cpp_digest=_content_digest(launch_cpp_text), + compile_config_digest=_content_digest(compile_config_text), link_config_digest=_content_digest(link_config_text), ) return artifacts.shared_library, launch_symbol diff --git a/ptodsl/ptodsl/_runtime/toolchain.py b/ptodsl/ptodsl/_runtime/toolchain.py index 938db50a71..df54f9fe74 100644 --- a/ptodsl/ptodsl/_runtime/toolchain.py +++ b/ptodsl/ptodsl/_runtime/toolchain.py @@ -15,6 +15,18 @@ def resolve_ptoas_binary() -> Path: + env_override = os.environ.get("PTOAS_BIN") + if env_override: + candidate = Path(env_override) + if candidate.is_file(): + return candidate + from_path = shutil.which(env_override) + if from_path: + return Path(from_path) + raise FileNotFoundError( + f"PTOAS_BIN is set but does not resolve to an existing executable: {env_override}" + ) + repo_root = Path(__file__).resolve().parents[4] candidates = [ repo_root / "build" / "tools" / "ptoas" / "ptoas", diff --git a/ptodsl/ptodsl/_tracing/module_builder.py b/ptodsl/ptodsl/_tracing/module_builder.py index f4108724c5..1625b07e9b 100644 --- a/ptodsl/ptodsl/_tracing/module_builder.py +++ b/ptodsl/ptodsl/_tracing/module_builder.py @@ -31,6 +31,7 @@ class KernelModuleSpec: function_name: str target_arch: str kernel_kind: str + kernel_kind_explicit: bool = False backend: str = "vpto" entry: bool = True mode: str = "auto" diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index c2390f1d63..8f81836996 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -13,7 +13,7 @@ from dataclasses import dataclass import hashlib -from .._diagnostics import inline_subkernel_value_escape_error +from .._diagnostics import inline_subkernel_value_escape_error, subkernel_kernel_kind_mismatch_error from .._kernel_signature import RuntimeScalarParameterSpec from .._ops import const from .._surface_values import unwrap_surface_value, wrap_like_surface_value @@ -239,16 +239,50 @@ def _create_subkernel_section_op(self, role: str): return None def _create_inline_subkernel_wrapper(self, role: str): - wrapper_op = self._create_subkernel_section_op(role) + wrapper_op = None + if self._subkernel_section_policy(role) != "function_kind": + wrapper_op = self._create_subkernel_section_op(role) if wrapper_op is None: wrapper_op = _pto.VecScopeOp() body_block = wrapper_op.body.blocks.append() return wrapper_op, body_block + def _subkernel_role_kernel_kind(self, role: str) -> str | None: + if role == "simd": + return "vector" + if role == "cube": + return "cube" + return None + + def _current_explicit_kernel_kind(self) -> str | None: + module_spec = self.current_function_module_spec + if not getattr(module_spec, "kernel_kind_explicit", False): + return None + kind = getattr(module_spec, "kernel_kind", None) + return kind if kind in {"cube", "vector"} else None + + def _subkernel_section_policy(self, role: str) -> str: + role_kind = self._subkernel_role_kernel_kind(role) + explicit_kind = self._current_explicit_kernel_kind() + if role_kind is None or explicit_kind is None: + return "section" + if explicit_kind != role_kind: + raise subkernel_kernel_kind_mismatch_error(role, explicit_kind) + return "function_kind" + def _subkernel_helper_attributes(self, role: str) -> tuple[tuple[str, object], ...]: attrs: list[tuple[str, object]] = [] if role in {"simd", "cube"}: attrs.append(("pto.ptodsl.subkernel_helper", StringAttr.get(role))) + if self._subkernel_section_policy(role) == "function_kind": + attrs.append( + ( + "pto.kernel_kind", + Attribute.parse( + f"#pto.kernel_kind<{self._subkernel_role_kernel_kind(role)}>" + ), + ) + ) if role == "simt": attrs.append(("pto.simt_entry", UnitAttr.get())) return tuple(attrs) @@ -275,6 +309,10 @@ def enter_subkernel_body(self, role: str, symbol_name: str, target: str): ) self._subkernel_stack.append(frame) try: + if self._subkernel_section_policy(role) == "function_kind": + yield frame + return + section_op = self._create_subkernel_section_op(role) if section_op is None: yield frame @@ -391,7 +429,8 @@ def _remap_captured_operands(self, root_ops, capture_mapping) -> None: def _outline_inline_subkernel(self, outline_frame: InlineSubkernelOutlineFrame) -> None: role = outline_frame.trace_frame.role - if role in {"simd", "cube"}: + section_policy = self._subkernel_section_policy(role) + if role in {"simd", "cube"} and section_policy != "function_kind": root_ops = (outline_frame.wrapper_op,) else: root_ops = tuple(outline_frame.body_block.operations) @@ -422,7 +461,7 @@ def _outline_inline_subkernel(self, outline_frame: InlineSubkernelOutlineFrame) terminator = func.ReturnOp([]) return_anchor = terminator.operation.opview - if role in {"simd", "cube"}: + if role in {"simd", "cube"} and section_policy != "function_kind": outline_frame.wrapper_op.move_before(return_anchor) outlined_roots = (outline_frame.wrapper_op,) else: diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 48cbbbf74a..557c4b3a2a 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -149,6 +149,25 @@ def host_vec_copy_explicit( pto.tile.store(o_tile, out) +@pto.jit(target="a5", mode="explicit") +def host_vec_copy_explicit_addr( + A_ptr: pto.ptr(pto.f32, "gm"), + O_ptr: pto.ptr(pto.f32, "gm"), + rows: pto.i32, + cols: pto.i32, + *, + BLOCK: pto.const_expr = 128, +): + a_view = pto.make_tensor_view(A_ptr, shape=[rows, cols], strides=[cols, 1]) + o_view = pto.make_tensor_view(O_ptr, shape=[rows, cols], strides=[cols, 1]) + a_tile = pto.alloc_tile(shape=[1, BLOCK], dtype=pto.f32, addr=0) + o_tile = pto.alloc_tile(shape=[1, BLOCK], dtype=pto.f32, addr=4096) + part = pto.partition_view(a_view, offsets=[0, 0], sizes=[rows, cols]) + out = pto.partition_view(o_view, offsets=[0, 0], sizes=[rows, cols]) + pto.tile.load(part, a_tile) + pto.tile.store(o_tile, out) + + @pto.jit(target="a5", backend="emitc") def host_vec_copy_emitc( A_ptr: pto.ptr(pto.f32, "gm"), @@ -668,6 +687,16 @@ def top_level_simd_probe(): SUBKERNEL_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth)) +@pto.simd +def explicit_vector_simd_probe(): + pto.pipe_barrier(pto.Pipe.ALL) + + +@pto.cube +def explicit_vector_cube_probe(): + pto.pipe_barrier(pto.Pipe.ALL) + + @pto.jit(target="a5") def shared_subkernel_lowering_probe(*, TRACE_TOKEN: pto.const_expr = 0): top_level_cube_probe() @@ -675,6 +704,22 @@ def shared_subkernel_lowering_probe(*, TRACE_TOKEN: pto.const_expr = 0): nested_simd_probe() +@pto.jit(target="a5", kernel_kind="vector") +def explicit_vector_calls_simd_probe(*, TRACE_TOKEN: pto.const_expr = 0): + explicit_vector_simd_probe() + + +@pto.jit(target="a5", kernel_kind="vector") +def explicit_vector_calls_cube_probe(*, TRACE_TOKEN: pto.const_expr = 0): + explicit_vector_cube_probe() + + +@pto.jit(target="a5", kernel_kind="vector") +def explicit_vector_inline_simd_probe(*, TRACE_TOKEN: pto.const_expr = 0): + with pto.simd(): + pto.pipe_barrier(pto.Pipe.ALL) + + @pto.jit(target="a5", mode="explicit") def inline_subkernel_scope_probe(*, TRACE_TOKEN: pto.const_expr = 0): session = current_session() @@ -3461,6 +3506,10 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): and helper_cache_signature[4] is False, "@pto.jit(entry=False) handles should expose an explicit, stable cache-signature protocol", ) + expect( + helper_cache_signature[7] == "vector" and helper_cache_signature[8] is False, + "default @pto.jit handles should keep vector as the effective kernel kind while recording that it was not explicit", + ) expect_raises( RuntimeError, kernel_module_return_probe.compile, @@ -3660,6 +3709,7 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): ) native_build_variants = ( ("pure-container", host_vec_copy.compile()), + ("explicit-level3-container", host_vec_copy_explicit_addr.compile()), ("same-backend-multi-child-container", kernel_module_compiled), ("mixed-backend-container", emitc_entry_calls_vpto_kernel_module_probe.compile()), ("source-auto", source_native_build_compiled), @@ -3748,14 +3798,14 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke f"{label} native build should forward the effective insert_sync policy to ptoas", ) expected_backend = compiled._module_spec.backend if compiled._module_spec.jit_source is not None else None - expected_pto_level = "level3" if compiled._module_spec.jit_source is not None and compiled._module_spec.mode == "explicit" else None + expected_pto_level = "level3" if compiled._module_spec.mode == "explicit" else None expect( observation["backend"] == expected_backend, f"{label} native build should only forward ptoas backend overrides for source-backed kernels", ) expect( observation["pto_level"] == expected_pto_level, - f"{label} native build should only map explicit mode to ptoas level3 for source-backed kernels", + f"{label} native build should derive the PTOAS level from the authored mode", ) expect( observation["mlir_text"] == compiled.mlir_text(), @@ -3797,7 +3847,7 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): ) expect( "--pto-level=level3" not in ptoas_cmd, - "native build should no longer reconstruct explicit mode through a global pto-level flag", + "native build should not pass a global pto-level flag by default", ) expect( "--enable-insert-sync" not in ptoas_cmd, @@ -3831,21 +3881,21 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): kernel_object, target_arch="a5", backend="vpto", - pto_level="level3", + pto_level=native_build_runtime._effective_pto_level(mode="explicit"), insert_sync=True, ) - expect(len(ptoas_cmds) == 1, "native build should issue exactly one ptoas command with source-backed overrides") - source_ptoas_cmd = ptoas_cmds[0] + expect(len(ptoas_cmds) == 1, "native build should issue exactly one ptoas command with explicit-mode PTOAS policy") + explicit_ptoas_cmd = ptoas_cmds[0] expect( - "--pto-backend=vpto" in source_ptoas_cmd, + "--pto-backend=vpto" in explicit_ptoas_cmd, "source-backed native build should pass the decorator backend to ptoas", ) expect( - "--pto-level=level3" in source_ptoas_cmd, - "source-backed explicit mode should pass --pto-level=level3 to ptoas", + "--pto-level=level3" in explicit_ptoas_cmd, + 'native build should pass --pto-level=level3 for mode="explicit"', ) expect( - "--enable-insert-sync" in source_ptoas_cmd, + "--enable-insert-sync" in explicit_ptoas_cmd, "source-backed native build should still pass explicit/effective insert-sync to ptoas", ) expect("valid=?" not in default_text, "default alloc_tile() should keep full static valid-shape when valid_shape= is omitted") @@ -4081,6 +4131,34 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): "outlined decorated helper bodies should still preserve their PTO unit sections", ) + explicit_vector_simd_text = explicit_vector_calls_simd_probe.compile(TRACE_TOKEN=1).mlir_text() + expect_parse_roundtrip_and_verify( + explicit_vector_simd_text, + "explicit vector jit calling simd subkernel specialization", + ) + expect( + "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_simd_text + and "pto.section.vector {" not in explicit_vector_simd_text, + "same-kind @pto.simd helpers inside explicit vector kernels should use function/kernel kind metadata without redundant sections", + ) + expect_raises( + RuntimeError, + lambda: explicit_vector_calls_cube_probe.compile(TRACE_TOKEN=1).mlir_text(), + "@pto.cube cannot be lowered inside an explicit @pto.jit(kernel_kind='vector')", + ) + explicit_vector_inline_simd_text = explicit_vector_inline_simd_probe.compile( + TRACE_TOKEN=1 + ).mlir_text() + expect_parse_roundtrip_and_verify( + explicit_vector_inline_simd_text, + "explicit vector jit calling inline simd specialization", + ) + expect( + "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_inline_simd_text + and "pto.section.vector {" not in explicit_vector_inline_simd_text, + "same-kind inline pto.simd() scopes inside explicit vector kernels should avoid redundant sections", + ) + INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.clear() inline_subkernel_scope_text = inline_subkernel_scope_probe.compile(TRACE_TOKEN=1).mlir_text() expect_parse_roundtrip_and_verify(inline_subkernel_scope_text, "inline subkernel scope specialization") diff --git a/ptodsl/tests/test_runtime_toolchain.py b/ptodsl/tests/test_runtime_toolchain.py new file mode 100644 index 0000000000..041c1aa258 --- /dev/null +++ b/ptodsl/tests/test_runtime_toolchain.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +import os +import tempfile +import unittest +from unittest import mock + +from pathlib import Path + +from ptodsl._runtime import toolchain + + +class ResolvePtoasBinaryTests(unittest.TestCase): + def test_env_override_wins_over_repo_default(self): + with tempfile.TemporaryDirectory() as temp_dir: + temp_root = Path(temp_dir) + env_ptoas = temp_root / "env-ptoas" + env_ptoas.write_text("", encoding="utf-8") + + repo_build_ptoas = temp_root / "build" / "tools" / "ptoas" / "ptoas" + repo_build_ptoas.parent.mkdir(parents=True, exist_ok=True) + repo_build_ptoas.write_text("", encoding="utf-8") + + fake_toolchain_file = temp_root / "repo" / "ptodsl" / "ptodsl" / "_runtime" / "toolchain.py" + fake_toolchain_file.parent.mkdir(parents=True, exist_ok=True) + fake_toolchain_file.write_text("", encoding="utf-8") + + with mock.patch.dict(os.environ, {"PTOAS_BIN": str(env_ptoas)}, clear=False), mock.patch.object( + toolchain, "__file__", str(fake_toolchain_file) + ): + resolved = toolchain.resolve_ptoas_binary() + + self.assertEqual(resolved, env_ptoas) + + def test_invalid_env_override_raises(self): + with tempfile.TemporaryDirectory() as temp_dir: + missing = Path(temp_dir) / "missing-ptoas" + with mock.patch.dict(os.environ, {"PTOAS_BIN": str(missing)}, clear=False): + with self.assertRaises(FileNotFoundError): + toolchain.resolve_ptoas_binary() + + +if __name__ == "__main__": + unittest.main() diff --git a/scripts/ptoas_env.sh b/scripts/ptoas_env.sh index ac5eaa68b9..a074f7f34a 100644 --- a/scripts/ptoas_env.sh +++ b/scripts/ptoas_env.sh @@ -134,7 +134,13 @@ echo "[ptoas_env] PTO_ISA_PATH=${PTO_ISA_PATH}" echo "[ptoas_env] ASCEND_HOME_PATH=${ASCEND_HOME_PATH}" echo "[ptoas_env] PATH/PYTHONPATH/LD_LIBRARY_PATH updated" -if [[ "${PTOAS_ENV_SKIP_SMOKE_TEST:-${CI:-0}}" != "1" ]]; then +_ptoas_env_skip_smoke="${PTOAS_ENV_SKIP_SMOKE_TEST:-0}" +if [[ -z "${PTOAS_ENV_SKIP_SMOKE_TEST:-}" && + ("${CI:-}" == "1" || "${CI:-}" == "true") ]]; then + _ptoas_env_skip_smoke=1 +fi + +if [[ "${_ptoas_env_skip_smoke}" != "1" ]]; then _ptoas_run_legacy_smoke_test fi @@ -142,3 +148,4 @@ unset -f _ptoas_prepend_path unset -f _ptoas_run_legacy_smoke_test unset _PTOAS_ENV_SCRIPT_DIR unset _PTOAS_REPO_DIR +unset _ptoas_env_skip_smoke diff --git a/scripts/sim_dsl.sh b/scripts/sim_dsl.sh index cd0bccbdf8..5632fb2e92 100755 --- a/scripts/sim_dsl.sh +++ b/scripts/sim_dsl.sh @@ -35,6 +35,8 @@ Environment: Keep the private staging directory after a successful sync. PTOAS_MSPROF_LOG_MODE=quiet|verbose Override the default simulator log rendering mode. + PYTHON_BIN Python executable used for the PTODSL example. + Defaults to python3. Examples: scripts/sim_dsl.sh ptodsl/examples/jit/tadd_launch.py @@ -177,12 +179,32 @@ ensure_private_dir "${PRIVATE_ROOT}" RUNTIME_OUTPUT_DIR="$(mktemp -d "${PRIVATE_ROOT}/${EXAMPLE_STEM}.XXXXXX")" chmod 700 "${RUNTIME_OUTPUT_DIR}" MSPROF_STDIO_LOG="${RUNTIME_OUTPUT_DIR}/msprof.stdout.log" +EXAMPLE_EXIT_CODE_FILE="${RUNTIME_OUTPUT_DIR}/example.exitcode" +EXAMPLE_LAUNCHER="${RUNTIME_OUTPUT_DIR}/run_example.sh" +PYTHON_BIN="${PTO_PYTHON_BIN:-${PYTHON_BIN:-python3}}" source "${ASCEND_HOME_PATH}/bin/setenv.bash" source "${REPO_ROOT}/scripts/ptoas_env.sh" export LD_LIBRARY_PATH="${SIM_LIB_DIR}:${LD_LIBRARY_PATH:-}" ulimit -n 65535 +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + die "PYTHON_BIN is not executable or not found on PATH: ${PYTHON_BIN}" +fi + +cat > "${EXAMPLE_LAUNCHER}" <<'EOF' +#!/usr/bin/env bash +set +e +"${PTOAS_SIM_DSL_PYTHON_BIN}" "${PTOAS_SIM_DSL_EXAMPLE_PATH}" "$@" +status=$? +printf '%s\n' "${status}" > "${PTOAS_SIM_DSL_EXIT_CODE_FILE}" +exit "${status}" +EOF +chmod 700 "${EXAMPLE_LAUNCHER}" +export PTOAS_SIM_DSL_PYTHON_BIN="${PYTHON_BIN}" +export PTOAS_SIM_DSL_EXAMPLE_PATH="${EXAMPLE_PATH}" +export PTOAS_SIM_DSL_EXIT_CODE_FILE="${EXAMPLE_EXIT_CODE_FILE}" + # msprof rejects group/other-writable working directories, so always launch # from a private directory and use an absolute path for the example script. cd "${HOME}" @@ -196,11 +218,30 @@ set +e msprof op simulator \ --soc-version="${SOC_VERSION}" \ --output="${RUNTIME_OUTPUT_DIR}" \ - python3 "${EXAMPLE_PATH}" "${EXAMPLE_ARGS[@]}" \ + "${EXAMPLE_LAUNCHER}" "${EXAMPLE_ARGS[@]}" \ > "${MSPROF_STDIO_LOG}" 2>&1 -STATUS=$? +MSPROF_STATUS=$? set -e +EXAMPLE_STATUS=0 +if [[ -f "${EXAMPLE_EXIT_CODE_FILE}" ]]; then + EXAMPLE_STATUS="$(< "${EXAMPLE_EXIT_CODE_FILE}")" + if [[ ! "${EXAMPLE_STATUS}" =~ ^[0-9]+$ ]]; then + log "invalid example exit code recorded in ${EXAMPLE_EXIT_CODE_FILE}: ${EXAMPLE_STATUS}" + EXAMPLE_STATUS=1 + fi +else + log "example exit code file was not produced: ${EXAMPLE_EXIT_CODE_FILE}" + EXAMPLE_STATUS=1 +fi + +STATUS=0 +if [[ ${MSPROF_STATUS} -ne 0 ]]; then + STATUS=${MSPROF_STATUS} +elif [[ ${EXAMPLE_STATUS} -ne 0 ]]; then + STATUS=${EXAMPLE_STATUS} +fi + print_msprof_log "${MSPROF_STDIO_LOG}" "${MSPROF_LOG_MODE}" "${STATUS}" SYNC_STATUS=0 diff --git a/test/dsl-st/cube_matrix_pipeline.py b/test/dsl-st/cube_matrix_pipeline.py index 74485690e5..f4836eac30 100644 --- a/test/dsl-st/cube_matrix_pipeline.py +++ b/test/dsl-st/cube_matrix_pipeline.py @@ -15,50 +15,19 @@ M = 16 K = 32 -N = 48 +N = 64 +ELEM_BYTES = 4 L1_A_ADDR = 0 L1_B_ADDR = 4096 -UB_O_ADDR = 0 L0A_ADDR = 0 L0B_ADDR = 0 L0C_ADDR = 0 -@pto.cube -def cube_gemm_tile(a_mat, b_mat, o_tile, a_l0a, b_l0b, o_acc): - m = a_mat.valid_shape[0] - k = a_mat.valid_shape[1] - n = b_mat.valid_shape[1] - - pto.mte_l1_l0a(a_mat.as_ptr(), a_l0a.as_ptr(), m, k) - pto.mte_l1_l0b(b_mat.as_ptr(), b_l0b.as_ptr(), k, n) - pto.set_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) - pto.wait_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) - pto.mad( - a_l0a.as_ptr(), - b_l0b.as_ptr(), - o_acc.as_ptr(), - m, - n, - k, - unit_flag=pto.MadUnitFlagMode.CHECK_ONLY, - sat=pto.SatMode.OFF, - ) - pto.set_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) - pto.wait_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) - pto.mte_l0c_ub( - o_acc.as_ptr(), - o_tile.as_ptr(), - m, - n, - n, - n, - ) - - @pto.jit( name="cube_matrix_pipeline_kernel", + kernel_kind="cube", target="a5", mode="explicit", insert_sync=False, @@ -68,20 +37,14 @@ def cube_matrix_pipeline_kernel( b_ptr: pto.ptr(pto.f32, "gm"), o_ptr: pto.ptr(pto.f32, "gm"), ): - a_view = pto.make_tensor_view(a_ptr, shape=[M, K], strides=[K, 1]) - b_view = pto.make_tensor_view(b_ptr, shape=[K, N], strides=[N, 1]) - o_view = pto.make_tensor_view(o_ptr, shape=[M, N], strides=[N, 1]) - - a_part = pto.partition_view(a_view, offsets=[0, 0], sizes=[M, K]) - b_part = pto.partition_view(b_view, offsets=[0, 0], sizes=[K, N]) - o_part = pto.partition_view(o_view, offsets=[0, 0], sizes=[M, N]) - a_mat = pto.alloc_tile( shape=[M, K], dtype=pto.f32, memory_space=pto.MemorySpace.MAT, addr=L1_A_ADDR, valid_shape=[M, K], + blayout="ColMajor", + slayout="RowMajor", ) b_mat = pto.alloc_tile( shape=[K, N], @@ -89,12 +52,8 @@ def cube_matrix_pipeline_kernel( memory_space=pto.MemorySpace.MAT, addr=L1_B_ADDR, valid_shape=[K, N], - ) - o_tile = pto.alloc_tile( - shape=[M, N], - dtype=pto.f32, - addr=UB_O_ADDR, - valid_shape=[M, N], + blayout="ColMajor", + slayout="RowMajor", ) a_l0a = pto.alloc_tile( shape=[M, K], @@ -102,6 +61,8 @@ def cube_matrix_pipeline_kernel( memory_space=pto.MemorySpace.LEFT, addr=L0A_ADDR, valid_shape=[M, K], + blayout="ColMajor", + slayout="RowMajor", ) b_l0b = pto.alloc_tile( shape=[K, N], @@ -109,6 +70,8 @@ def cube_matrix_pipeline_kernel( memory_space=pto.MemorySpace.RIGHT, addr=L0B_ADDR, valid_shape=[K, N], + blayout="RowMajor", + slayout="ColMajor", ) o_acc = pto.alloc_tile( shape=[M, N], @@ -116,16 +79,58 @@ def cube_matrix_pipeline_kernel( memory_space=pto.MemorySpace.ACC, addr=L0C_ADDR, valid_shape=[M, N], + blayout="ColMajor", + slayout="RowMajor", + fractal_size=1024, ) - pto.tile.load(a_part, a_mat) - pto.tile.load(b_part, b_mat) + a_l1_ptr = pto.castptr(pto.ui64(L1_A_ADDR), pto.ptr(pto.f32, "mat")) + b_l1_ptr = pto.castptr(pto.ui64(L1_B_ADDR), pto.ptr(pto.f32, "mat")) + + pto.mte_gm_l1_frac( + a_ptr, + a_l1_ptr, + pto.FractalMode.ND2NZ, + shape=(M, K), + src_layout=(K * ELEM_BYTES,), + dst_group=(1, 1, M, 0), + ctrl=(0, False), + ) pto.set_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=0) pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=0) - cube_gemm_tile(a_mat, b_mat, o_tile, a_l0a, b_l0b, o_acc) - pto.set_flag(pto.Pipe.FIX, pto.Pipe.MTE3, event_id=2) - pto.wait_flag(pto.Pipe.FIX, pto.Pipe.MTE3, event_id=2) - pto.tile.store(o_tile, o_part) + pto.mte_l1_l0a(a_l1_ptr, a_l0a.as_ptr(), M, K) + + pto.mte_gm_l1_frac( + b_ptr, + b_l1_ptr, + pto.FractalMode.ND2NZ, + shape=(K, N), + src_layout=(N * ELEM_BYTES,), + dst_group=(1, 1, K, 0), + ctrl=(0, False), + ) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=1) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=1) + pto.mte_l1_l0b(b_l1_ptr, b_l0b.as_ptr(), K, N, transpose=True) + + pto.set_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) + pto.wait_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) + pto.tile.matmul(a_l0a, b_l0b, o_acc) + + pto.set_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) + pto.wait_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) + pto.mte_l0c_gm( + o_acc.as_ptr(), + o_ptr, + M, + N, + M, + N, + 0, + 0, + layout="nz2nd", + ) + pto.pipe_barrier(pto.Pipe.ALL) def make_inputs(): diff --git a/test/dsl-st/gemv_mx_pipeline.py b/test/dsl-st/gemv_mx_pipeline.py index 41c64531ff..985ab94768 100644 --- a/test/dsl-st/gemv_mx_pipeline.py +++ b/test/dsl-st/gemv_mx_pipeline.py @@ -11,8 +11,6 @@ from common import assert_close, auto_main from ptodsl import pto -from ptodsl._surface_values import unwrap_surface_value -from mlir.dialects import pto as _pto M = 1 @@ -171,17 +169,6 @@ def _alloc_common_tiles(): return lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile -def _bind_mx_scale_tiles(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile): - _pto.TGetScaleAddrOp( - unwrap_surface_value(lhs_tile), - unwrap_surface_value(lhs_scale_tile), - ) - _pto.TGetScaleAddrOp( - unwrap_surface_value(rhs_tile), - unwrap_surface_value(rhs_scale_tile), - ) - - def _alloc_bias_tile(): return pto.alloc_tile( shape=[M, N_STORAGE], @@ -265,7 +252,6 @@ def gemv_mx_fp8_pipeline_kernel( ): lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile = _alloc_common_tiles() _stage_fp8_tiles(a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, lhs_tile, rhs_tile) - _bind_mx_scale_tiles(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile) pto.tile.gemv_mx(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile) _writeback_output(dst_tile, out_ptr) @@ -286,7 +272,6 @@ def gemv_mx_acc_fp8_pipeline_kernel( ): lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile = _alloc_common_tiles() _stage_fp8_tiles(a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, lhs_tile, rhs_tile) - _bind_mx_scale_tiles(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile) pto.tile.gemv_mx(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile) pto.tile.gemv_mx_acc(dst_tile, lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, dst_tile) _writeback_output(dst_tile, out_ptr) @@ -319,7 +304,6 @@ def gemv_mx_bias_fp8_pipeline_kernel( bias_ptr=bias_ptr, bias_tile=bias_tile, ) - _bind_mx_scale_tiles(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile) pto.tile.gemv_mx_bias(lhs_tile, lhs_scale_tile, rhs_tile, rhs_scale_tile, bias_tile, dst_tile) _writeback_output(dst_tile, out_ptr) diff --git a/test/dsl-st/predicate_pack.py b/test/dsl-st/predicate_pack.py index b170f8b714..f39252efc8 100644 --- a/test/dsl-st/predicate_pack.py +++ b/test/dsl-st/predicate_pack.py @@ -76,37 +76,38 @@ def predicate_pack_part_kernel( addr=1024, valid_shape=[ROWS, ROW_BYTES], ) + src_ptr = pto.castptr(pto.ui64(0), pto.ptr(pto.ui8, "ub")) + dst_ptr = pto.castptr(pto.ui64(1024), pto.ptr(pto.ui8, "ub")) pto.tile.load(inp_part, src_tile) pto.tile.load(out_part, dst_tile) pto.set_flag("MTE2", "V", event_id=0) pto.wait_flag("MTE2", "V", event_id=0) - with pto.simd(): - seed = pto.pset_b8(pto.MaskPattern.ALL) - src = pto.vlds(src_tile[0, 0:]) - active_b8 = pto.vcmp(src, src, seed, pto.CmpMode.EQ) - active = pto.pbitcast(active_b8, pto.mask_b32) - - packed_lo_same = pto.ppack(active, "LOWER") - unpacked_lo_same = pto.punpack(packed_lo_same, "LOWER") - packed_lo_b16 = pto.ppack(active, "LOWER", to_type=pto.mask_b16) - unpacked_lo_b32 = pto.punpack(packed_lo_b16, "LOWER", to_type=pto.mask_b32) - - packed_hi_same = pto.ppack(active, "HIGHER") - unpacked_hi_same = pto.punpack(packed_hi_same, "HIGHER") - packed_hi_b16 = pto.ppack(active, "HIGHER", to_type=pto.mask_b16) - unpacked_hi_b32 = pto.punpack(packed_hi_b16, "HIGHER", to_type=pto.mask_b32) - - pto.psts(active, dst_tile.as_ptr(), 0, dist="NORM") - pto.psts(packed_lo_same, dst_tile.as_ptr(), ROW_BYTES, dist="NORM") - pto.psts(unpacked_lo_same, dst_tile.as_ptr(), ROW_BYTES * 2, dist="NORM") - pto.psts(packed_lo_b16, dst_tile.as_ptr(), ROW_BYTES * 3, dist="NORM") - pto.psts(unpacked_lo_b32, dst_tile.as_ptr(), ROW_BYTES * 4, dist="NORM") - pto.psts(packed_hi_same, dst_tile.as_ptr(), ROW_BYTES * 5, dist="NORM") - pto.psts(unpacked_hi_same, dst_tile.as_ptr(), ROW_BYTES * 6, dist="NORM") - pto.psts(packed_hi_b16, dst_tile.as_ptr(), ROW_BYTES * 7, dist="NORM") - pto.psts(unpacked_hi_b32, dst_tile.as_ptr(), ROW_BYTES * 8, dist="NORM") + seed = pto.pset_b8(pto.MaskPattern.ALL) + src = pto.vlds(src_ptr, 0, pto.vreg_type(256, pto.ui8)) + active_b8 = pto.vcmp(src, src, seed, pto.CmpMode.EQ) + active = pto.pbitcast(active_b8, pto.mask_b32) + + packed_lo_same = pto.ppack(active, "LOWER") + unpacked_lo_same = pto.punpack(packed_lo_same, "LOWER") + packed_lo_b16 = pto.ppack(active, "LOWER", to_type=pto.mask_b16) + unpacked_lo_b32 = pto.punpack(packed_lo_b16, "LOWER", to_type=pto.mask_b32) + + packed_hi_same = pto.ppack(active, "HIGHER") + unpacked_hi_same = pto.punpack(packed_hi_same, "HIGHER") + packed_hi_b16 = pto.ppack(active, "HIGHER", to_type=pto.mask_b16) + unpacked_hi_b32 = pto.punpack(packed_hi_b16, "HIGHER", to_type=pto.mask_b32) + + pto.psts(active, dst_ptr, 0, dist="NORM") + pto.psts(packed_lo_same, dst_ptr, ROW_BYTES, dist="NORM") + pto.psts(unpacked_lo_same, dst_ptr, ROW_BYTES * 2, dist="NORM") + pto.psts(packed_lo_b16, dst_ptr, ROW_BYTES * 3, dist="NORM") + pto.psts(unpacked_lo_b32, dst_ptr, ROW_BYTES * 4, dist="NORM") + pto.psts(packed_hi_same, dst_ptr, ROW_BYTES * 5, dist="NORM") + pto.psts(unpacked_hi_same, dst_ptr, ROW_BYTES * 6, dist="NORM") + pto.psts(packed_hi_b16, dst_ptr, ROW_BYTES * 7, dist="NORM") + pto.psts(unpacked_hi_b32, dst_ptr, ROW_BYTES * 8, dist="NORM") pto.set_flag("V", "MTE3", event_id=0) pto.wait_flag("V", "MTE3", event_id=0) diff --git a/test/lit/vpto/fold_tile_buf_intrinsics_dead_view_cleanup.pto b/test/lit/vpto/fold_tile_buf_intrinsics_dead_view_cleanup.pto new file mode 100644 index 0000000000..10cc6a777e --- /dev/null +++ b/test/lit/vpto/fold_tile_buf_intrinsics_dead_view_cleanup.pto @@ -0,0 +1,45 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-fold-tile-buf-intrinsics %s -o /dev/null 2>&1 | FileCheck %s + +// CHECK-LABEL: func.func @dead_view_cleanup +// CHECK-NOT: memref.reinterpret_cast %arg0 +// CHECK-NOT: pto.make_tensor_view +// CHECK-NOT: pto.partition_view +// CHECK: pto.vadd +// CHECK: return + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @dead_view_cleanup(%src: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %a = pto.alloc_tile + : !pto.tile_buf + %b = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + pto.tadd ins(%a, %b : !pto.tile_buf, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) + + %view = pto.make_tensor_view %src, shape = [%c16, %c16], strides = [%c16, %c1] + : !pto.tensor_view + %part = pto.partition_view %view, offsets = [%c0, %c0], sizes = [%c16, %c16] + : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> + return + } +} diff --git a/test/tilelib-st/README.md b/test/tilelib-st/README.md new file mode 100644 index 0000000000..43cda3d89f --- /dev/null +++ b/test/tilelib-st/README.md @@ -0,0 +1,361 @@ +# TileLib ST + +`test/tilelib-st` contains PTODSL TileLib system tests that compile and run +operator-level TileLib kernels against a torch_npu runtime or the msprof +simulator wrapper. + +## Layout + +- Put A5 cases under `test/tilelib-st/a5/`. +- Put shared test helpers in `test/tilelib-st/common.py`. +- Keep one operator family per directory, for example `a5/tadd/` or + `a5/tmatmul/`. +- Put the executable case module in that directory's `case.py`. +- Each `case.py` file must define a non-empty `CASES` list. + +## Writing Cases + +Use `golden_output_case(...)` for the common "inputs + output + golden compare" +shape. Prefer this auto-mode style for simple vector TileLib cases: + +```python +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +ROWS = 16 +COLS = 64 + + +@pto.jit(name="my_tadd_f32_16x64", target="a5") +def _kernel( + a_ptr: pto.ptr(pto.f32, "gm"), + b_ptr: pto.ptr(pto.f32, "gm"), + out_ptr: pto.ptr(pto.f32, "gm"), +): + a_view = pto.make_tensor_view(a_ptr, shape=[ROWS, COLS], strides=[COLS, 1]) + b_view = pto.make_tensor_view(b_ptr, shape=[ROWS, COLS], strides=[COLS, 1]) + out_view = pto.make_tensor_view(out_ptr, shape=[ROWS, COLS], strides=[COLS, 1]) + + a_tile = pto.alloc_tile(shape=[ROWS, COLS], dtype=pto.f32) + b_tile = pto.alloc_tile(shape=[ROWS, COLS], dtype=pto.f32) + out_tile = pto.alloc_tile(shape=[ROWS, COLS], dtype=pto.f32) + + pto.tile.load(a_view, a_tile) + pto.tile.load(b_view, b_tile) + pto.tile.add(a_tile, b_tile, out_tile) + pto.tile.store(out_tile, out_view) + + +def _make_inputs(): + rng = np.random.default_rng(0xA5) + a = rng.uniform(-1.0, 1.0, size=(ROWS, COLS)).astype(np.float32) + b = rng.uniform(-1.0, 1.0, size=(ROWS, COLS)).astype(np.float32) + return [a, b] + + +def _make_expected(a, b): + return (a + b).astype(np.float32) + + +CASES = [ + golden_output_case( + "my_tadd_f32_16x64", + _kernel, + inputs=_make_inputs, + expected=_make_expected, + rtol=1e-6, + atol=1e-6, + ), +] + + +auto_main(globals()) +``` + +Key conventions: + +- Use `@pto.jit(target="a5")` for auto-mode vector tile-op cases. +- Build `TensorView` objects with `pto.make_tensor_view(...)`. +- Pass `TensorView` directly to `pto.tile.load/store`; it will infer the + full-tile partition from tile metadata. +- Allocate tiles with `pto.alloc_tile(shape=..., dtype=...)` and omit `addr=`. +- Put host-side inputs in `inputs=...` and NumPy golden logic in `expected=...`. +- Keep case names unique across the whole discovered tree. + +For multiple cases in one suite, create one `@pto.jit` kernel per static shape +or layout variant in that suite's `case.py`, then append one +`golden_output_case(...)` entry for each case: + +```python +CASE_SPECS = [ + ("f32_16x64", 16, 64), + ("f32_32x32", 32, 32), +] + +_kernels = {} +for _name, _rows, _cols in CASE_SPECS: + + def _make(rows=_rows, cols=_cols, kernel_name=f"my_op_{_name}"): + @pto.jit(name=kernel_name, target="a5") + def _kernel(src_ptr: pto.ptr(pto.f32, "gm"), out_ptr: pto.ptr(pto.f32, "gm")): + src_view = pto.make_tensor_view(src_ptr, shape=[rows, cols], strides=[cols, 1]) + out_view = pto.make_tensor_view(out_ptr, shape=[rows, cols], strides=[cols, 1]) + src_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) + out_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) + pto.tile.load(src_view, src_tile) + pto.tile.abs(src_tile, out_tile) + pto.tile.store(out_tile, out_view) + + return _kernel + + _kernels[_name] = _make() + + +CASES = [ + golden_output_case( + "my_op_" + name, + _kernels[name], + inputs=lambda rows=rows, cols=cols: [np.ones((rows, cols), dtype=np.float32)], + expected=lambda src: np.abs(src).astype(np.float32), + ) + for name, rows, cols in CASE_SPECS +] +``` + +When a case needs a non-default tile layout, keep the same auto-mode structure +and pass layout metadata to `alloc_tile(...)`, for example: + +```python +tile = pto.alloc_tile(shape=[16, 64], dtype=pto.f32, blayout="ColMajor") +view = pto.make_tensor_view(ptr, shape=[16, 64], strides=[1, 16]) +``` + +Prefer automatic tile address allocation by omitting `addr=` from +`pto.alloc_tile(...)`. Use explicit tile addresses only when a case is +intentionally validating address-sensitive behavior or needs to mirror an +existing hand-authored ST exactly. Prefer auto-mode TileLib authoring for +simple vector tile-op cases; use explicit mode only when the case needs manual +sync or explicit low-level DMA/cube movement surfaces. For example, cube cases +that call `pto.mte_gm_l1_frac`, `pto.mte_l1_l0a`, `pto.mte_l1_l0b`, or +`pto.mte_l0c_gm` must use explicit mode because those APIs are explicit-only. + +## Running + +List cases: + +```bash +python3 test/tilelib-st/run_tilelib_st.py test/tilelib-st/a5 --list +``` + +Run every A5 suite: + +```bash +python3 test/tilelib-st/run_tilelib_st.py test/tilelib-st/a5 +``` + +Run a single suite directly: + +```bash +python3 test/tilelib-st/a5/tadd/case.py +``` + +Run the A5 directory through the simulator wrapper: + +```bash +ASCEND_HOME_PATH=/path/to/cann \ +PTOAS_BIN=/path/to/ptoas \ +PYTHON_BIN=/path/to/python-with-torch-npu \ +scripts/sim_dsl.sh test/tilelib-st/run_tilelib_st.py -- test/tilelib-st/a5 +``` + +The runtime Python must provide `torch`, `torch_npu`, `numpy`, PTODSL, and the +matching MLIR PTO Python bindings. + +## Simulator Runbook + +Use this flow when bringing up `test/tilelib-st` on a simulator host from a +fresh checkout. The commands assume Linux, a VPTO-enabled LLVM build, CANN with +`msprof op simulator`, and a Python environment that can import `torch` and +`torch_npu`. + +Prerequisites: + +- CANN is installed and provides `msprof op simulator`. +- VPTO LLVM/MLIR has already been built by following the repository root + `README.md` build guide. Generic upstream LLVM is not enough. +- The runtime Python can import `torch`, `torch_npu`, `numpy`, `pybind11`, + `nanobind`, and `yaml`. +- The Python ABI used to build PTOAS Python bindings matches the runtime + Python used to launch the ST cases. + +Use the LLVM/MLIR version documented in the repository root `README.md`. At the +time of writing that is LLVM 21 from the VPTO branch +`vpto-dev/llvm-project:feature-vpto-llvm21`. If unsure, check +`.github/workflows/ci_sim.yml` and use its current `LLVM_REPO` / `LLVM_REF`. +PTOAS CMake also verifies the LLVM major version and will reject an +incompatible LLVM build. + +Prepare the main paths: + +```bash +cd /path/to/PTOAS + +export PTOAS_REPO="$PWD" +export PTOAS_BUILD="$PTOAS_REPO/build-sim" +export ASCEND_HOME_PATH=/path/to/cann +export LLVM_BUILD_DIR=/path/to/vpto-llvm/build +export MLIR_PYTHON_ROOT="$LLVM_BUILD_DIR/tools/mlir/python_packages/mlir_core" +export PYTHON_BIN=/path/to/python-with-torch-npu +export PTO_PYTHON_BIN="$PYTHON_BIN" +export PTOAS_ENV_SKIP_SMOKE_TEST=1 +export TORCH_DEVICE_BACKEND_AUTOLOAD=0 +``` + +Source CANN before importing `torch_npu` or launching the simulator: + +```bash +source "$ASCEND_HOME_PATH/set_env.sh" >/dev/null 2>&1 || \ +source "$ASCEND_HOME_PATH/bin/setenv.bash" +``` + +Check the simulator, LLVM, and Python prerequisites: + +```bash +command -v msprof +test -d "$LLVM_BUILD_DIR/lib/cmake/llvm" +test -d "$LLVM_BUILD_DIR/lib/cmake/mlir" +test -d "$MLIR_PYTHON_ROOT" + +"$PYTHON_BIN" - <<'PY' +import nanobind +import numpy +import pybind11 +import torch +import torch_npu # noqa: F401 +import yaml + +print("nanobind", getattr(nanobind, "__version__", "unknown")) +print("numpy", numpy.__version__) +print("pybind11", pybind11.__version__) +print("torch", torch.__version__) +print("torch_npu", getattr(torch_npu, "__version__", "unknown")) +PY +``` + +If the lightweight build dependencies are missing, install them into the same +runtime Python. Do not install `torch` or `torch_npu` blindly; use a known +compatible prebuilt environment for those packages. + +```bash +"$PYTHON_BIN" -m pip install 'pybind11<3' nanobind numpy ml-dtypes PyYAML +``` + +Configure PTOAS if `build-sim` does not exist yet. The Python used to build +PTOAS Python bindings must have the same Python ABI as the runtime Python that +will run the ST cases: + +```bash +PYBIND11_CMAKE_DIR="$("$PYTHON_BIN" -m pybind11 --cmakedir)" +NANOBIND_CMAKE_DIR="$("$PYTHON_BIN" -m nanobind --cmake_dir)" + +cmake -S "$PTOAS_REPO" -B "$PTOAS_BUILD" -G Ninja \ + -DLLVM_DIR="$LLVM_BUILD_DIR/lib/cmake/llvm" \ + -DMLIR_DIR="$LLVM_BUILD_DIR/lib/cmake/mlir" \ + -DPython3_EXECUTABLE="$PYTHON_BIN" \ + -DPython3_FIND_STRATEGY=LOCATION \ + -Dpybind11_DIR="$PYBIND11_CMAKE_DIR" \ + -Dnanobind_DIR="$NANOBIND_CMAKE_DIR" \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DMLIR_PYTHON_PACKAGE_DIR="$MLIR_PYTHON_ROOT" \ + -DPTO_ENABLE_PYTHON_BINDING=ON \ + -DPTOAS_ENABLE_WERROR=OFF \ + -DBUILD_TESTING=ON +``` + +Build the compiler and PTO Python bindings: + +```bash +ninja -C "$PTOAS_BUILD" ptoas PTOPythonModules +``` + +Export the runtime lookup paths. `PTOAS_BIN` is useful, but PTODSL native build +also resolves `ptoas` through `PATH`, so keep the freshly built binary first: + +```bash +export PTOAS_BIN="$PTOAS_BUILD/tools/ptoas/ptoas" +export PATH="$PTOAS_BUILD/tools/ptoas:$PATH" +export PYTHONPATH="$PTOAS_BUILD/python:$PTOAS_REPO/ptodsl:$MLIR_PYTHON_ROOT:${PYTHONPATH:-}" +export LD_LIBRARY_PATH="$LLVM_BUILD_DIR/lib:$PTOAS_BUILD/lib:${LD_LIBRARY_PATH:-}" +``` + +Check that PTODSL and the PTO MLIR Python dialect import from the intended +build: + +```bash +"$PYTHON_BIN" - <<'PY' +from ptodsl import pto # noqa: F401 +from mlir.dialects import pto as _pto # noqa: F401 + +print("ptodsl imports ok") +PY +``` + +List the cases before running the simulator: + +```bash +"$PYTHON_BIN" test/tilelib-st/run_tilelib_st.py test/tilelib-st/a5 --list +``` + +Run all A5 TileLib ST cases through `msprof op simulator`: + +```bash +scripts/sim_dsl.sh --soc-version Ascend950PR_9599 \ + test/tilelib-st/run_tilelib_st.py -- test/tilelib-st/a5 \ + 2>&1 | tee /tmp/tilelib-st-a5-sim.log +``` + +Run one suite when debugging: + +```bash +scripts/sim_dsl.sh --soc-version Ascend950PR_9599 \ + test/tilelib-st/a5/tadd/case.py \ + 2>&1 | tee /tmp/tilelib-st-tadd-sim.log +``` + +The expected successful ending is: + +```text +PASS ... +All cases passed. +``` + +Example path layout: + +```bash +export ASCEND_HOME_PATH=/opt/ascend/cann +export LLVM_BUILD_DIR=/opt/vpto-llvm/build +export MLIR_PYTHON_ROOT="$LLVM_BUILD_DIR/tools/mlir/python_packages/mlir_core" +export PYTHON_BIN=/opt/ptodsl-runtime/bin/python +``` + +Common failures: + +- `libhccl.so: cannot open shared object file`: source the CANN environment + before importing `torch_npu`. +- `simulator library directory not found`: check + `$ASCEND_HOME_PATH/tools/simulator//lib` and pass the matching + `--soc-version`. +- `No TileLib ST cases discovered`: pass the case root through + `scripts/sim_dsl.sh ... run_tilelib_st.py -- test/tilelib-st/a5`, keeping + the `--` separator. +- Stale compiler behavior: rebuild `ptoas PTOPythonModules`, put + `$PTOAS_BUILD/tools/ptoas` first in `PATH`, and verify `$PTOAS_BIN --version`. diff --git a/test/tilelib-st/a5/__main__.py b/test/tilelib-st/a5/__main__.py new file mode 100644 index 0000000000..0501304652 --- /dev/null +++ b/test/tilelib-st/a5/__main__.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Directory runner for the A5 TileLib operator-level ST cases. +# +# Each operator suite (e.g. tadd, tmatmul, ...) is a subdirectory with a +# ``case.py`` module that authors its kernels with PTODSL and builds its CASES +# list through the helpers in the parent ``test/tilelib-st/common.py`` module. +# Running this directory discovers every nested ``case.py`` module and executes +# the cases against the torch_npu / simulator runtime. +# +# See test/tilelib-st/README.md for the authoring conventions. + +from pathlib import Path +import sys + + +if __package__ in {None, ""}: + # common.py lives one level up, in test/tilelib-st/. + sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from common import run_discovered_cases + + +if __name__ == "__main__": + raise SystemExit(run_discovered_cases(Path(__file__).resolve().parent)) diff --git a/test/tilelib-st/a5/tadd/case.py b/test/tilelib-st/a5/tadd/case.py new file mode 100644 index 0000000000..32cc5ab38d --- /dev/null +++ b/test/tilelib-st/a5/tadd/case.py @@ -0,0 +1,102 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# PTODSL rewrite of test/tilelang_st/npu/a5/src/st/testcase/tadd. +# +# This case intentionally uses PTODSL auto mode as the vector TileLib pilot: +# tile addresses, load/store partitions, and sync insertion are left to PTOAS. + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +# Each case is (name, shape). Both use fully-valid f32 tiles, matching the +# original tadd cases "f32_16x64" and "f32_32x32". +CASE_SHAPES = [ + ("f32_16x64", (16, 64)), + ("f32_32x32", (32, 32)), +] + +def _tadd_body(a_ptr, b_ptr, c_ptr, *, rows, cols): + """Shared kernel body for the two tadd cases.""" + + a_view = pto.make_tensor_view(a_ptr, shape=[rows, cols], strides=[cols, 1]) + b_view = pto.make_tensor_view(b_ptr, shape=[rows, cols], strides=[cols, 1]) + c_view = pto.make_tensor_view(c_ptr, shape=[rows, cols], strides=[cols, 1]) + + a_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) + b_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) + c_tile = pto.alloc_tile(shape=[rows, cols], dtype=pto.f32) + + pto.tile.load(a_view, a_tile) + pto.tile.load(b_view, b_tile) + pto.tile.add(a_tile, b_tile, c_tile) + pto.tile.store(c_tile, c_view) + + +# One decorated kernel per case, each binding a static shape at definition time +# (mirroring the per-case funcs in tadd.pto). +_tadd_kernels = {} +for _name, _shape in CASE_SHAPES: + _r, _c = _shape + + def _make(r=_r, c=_c, kernel_name=f"tadd_{_name}"): + @pto.jit( + name=kernel_name, + target="a5", + ) + def _kernel( + a_ptr: pto.ptr(pto.f32, "gm"), + b_ptr: pto.ptr(pto.f32, "gm"), + c_ptr: pto.ptr(pto.f32, "gm"), + ): + _tadd_body(a_ptr, b_ptr, c_ptr, rows=r, cols=c) + + return _kernel + + _tadd_kernels[_name] = _make() + + +def _make_inputs(name, shape): + # Deterministic per-case seed, mirroring st_common.setup_case_rng which uses + # crc32(name). Original value range was randint(1, 10). + import zlib + np.random.seed(zlib.crc32(name.encode("utf-8")) & 0xFFFFFFFF) + a = np.random.randint(1, 10, size=shape).astype(np.float32) + b = np.random.randint(1, 10, size=shape).astype(np.float32) + return [a, b] + + +def _make_expected(a, b): + return (a + b).astype(np.float32) + + +CASES = [] +for _name, _shape in CASE_SHAPES: + CASES.append( + golden_output_case( + "tadd_" + _name, + _tadd_kernels[_name], + inputs=lambda _name=_name, _shape=_shape: _make_inputs(_name, _shape), + expected=_make_expected, + rtol=1e-6, + atol=1e-6, + ) + ) + + +auto_main(globals()) diff --git a/test/tilelib-st/a5/tcolexpand/case.py b/test/tilelib-st/a5/tcolexpand/case.py new file mode 100644 index 0000000000..ca395ee158 --- /dev/null +++ b/test/tilelib-st/a5/tcolexpand/case.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Minimal PTODSL broadcast pilot for A5: +# tload(src) + tcolexpand(src)->dst + tstore(dst) + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +SRC_ROWS = 1 +DST_ROWS = 8 +COLS = 128 + + +@pto.jit( + name="tcolexpand_f32_1x8x128", + target="a5", +) +def _tcolexpand_kernel( + src_ptr: pto.ptr(pto.f32, "gm"), + dst_ptr: pto.ptr(pto.f32, "gm"), +): + src_view = pto.make_tensor_view( + src_ptr, + shape=[SRC_ROWS, COLS], + strides=[COLS, 1], + ) + dst_view = pto.make_tensor_view( + dst_ptr, + shape=[DST_ROWS, COLS], + strides=[COLS, 1], + ) + + src_tile = pto.alloc_tile(shape=[SRC_ROWS, COLS], dtype=pto.f32) + dst_tile = pto.alloc_tile(shape=[DST_ROWS, COLS], dtype=pto.f32) + + pto.tile.load(src_view, src_tile) + pto.tile.colexpand(src_tile, dst_tile) + pto.tile.store(dst_tile, dst_view) + + +def _make_input(): + rng = np.random.default_rng(0xC01E0A5) + return rng.uniform(-2.0, 2.0, size=(SRC_ROWS, COLS)).astype(np.float32) + + +def _make_expected(src): + return np.repeat(src, DST_ROWS, axis=0).astype(np.float32) + + +CASES = [ + golden_output_case( + "tcolexpand_f32_1x8x128", + _tcolexpand_kernel, + inputs=lambda: [_make_input()], + expected=_make_expected, + rtol=1e-6, + atol=1e-6, + ), +] + + +auto_main(globals()) diff --git a/test/tilelib-st/a5/tcolsum/case.py b/test/tilelib-st/a5/tcolsum/case.py new file mode 100644 index 0000000000..db0f08768c --- /dev/null +++ b/test/tilelib-st/a5/tcolsum/case.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Minimal PTODSL reduction pilot for A5: +# tload(src) + tcolsum(src)->dst + tstore(dst) + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +ROWS = 16 +COLS = 128 + + +@pto.jit( + name="tcolsum_f32_16x128", + target="a5", +) +def _tcolsum_kernel( + src_ptr: pto.ptr(pto.f32, "gm"), + dst_ptr: pto.ptr(pto.f32, "gm"), +): + src_view = pto.make_tensor_view( + src_ptr, + shape=[ROWS, COLS], + strides=[COLS, 1], + ) + dst_view = pto.make_tensor_view( + dst_ptr, + shape=[1, COLS], + strides=[COLS, 1], + ) + + src_tile = pto.alloc_tile(shape=[ROWS, COLS], dtype=pto.f32) + dst_tile = pto.alloc_tile(shape=[1, COLS], dtype=pto.f32) + + pto.tile.load(src_view, src_tile) + pto.tile.colsum(src_tile, dst_tile) + pto.tile.store(dst_tile, dst_view) + + +def _make_input(): + rng = np.random.default_rng(0xC01A5EED) + return rng.uniform(-3.0, 3.0, size=(ROWS, COLS)).astype(np.float32) + + +def _make_expected(src): + return np.sum(src, axis=0, keepdims=True, dtype=np.float32) + + +CASES = [ + golden_output_case( + "tcolsum_f32_16x128", + _tcolsum_kernel, + inputs=lambda: [_make_input()], + expected=_make_expected, + rtol=1e-5, + atol=1e-5, + ), +] + + +auto_main(globals()) diff --git a/test/tilelib-st/a5/tload_store/case.py b/test/tilelib-st/a5/tload_store/case.py new file mode 100644 index 0000000000..379a7e85a9 --- /dev/null +++ b/test/tilelib-st/a5/tload_store/case.py @@ -0,0 +1,129 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# PTODSL rewrite of the minimal GM -> tile -> GM coverage from +# test/tilelang_st/npu/a5/src/st/testcase/tload/tload.pto. +# +# Start with two static f32 round-trips: +# 1. ND / row-major +# 2. DN / col-major +# These are the smallest data-movement cases needed to validate that PTODSL can +# drive tload/tstore on A5 without the tilelang_st harness. + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +CASE_SPECS = [ + { + "case_name": "nd_f32_16x64", + "kernel_name": "tload_store_nd_f32_16x64", + "shape": (16, 64), + "view_strides": None, + "tile_kwargs": {}, + }, + { + "case_name": "dn_f32_16x64", + "kernel_name": "tload_store_dn_f32_16x64", + "shape": (16, 64), + "view_strides": None, + "tile_kwargs": {"blayout": "ColMajor"}, + }, +] + +def _roundtrip_body(src_ptr, dst_ptr, *, rows, cols, view_strides=None, tile_kwargs=None): + total = rows * cols + if view_strides is None: + view_strides = [total, total, total, cols, 1] + + if view_strides is not None and len(view_strides) == 5: + view_strides = view_strides[-2:] + + src_view = pto.make_tensor_view(src_ptr, shape=[rows, cols], strides=view_strides) + dst_view = pto.make_tensor_view(dst_ptr, shape=[rows, cols], strides=view_strides) + + tile = pto.alloc_tile( + shape=[rows, cols], + dtype=pto.f32, + **(tile_kwargs or {}), + ) + + pto.tile.load(src_view, tile) + pto.tile.store(tile, dst_view) + + +_tload_store_kernels = {} +for _spec in CASE_SPECS: + _rows, _cols = _spec["shape"] + _view_strides = _spec["view_strides"] + if _view_strides is None and _spec["tile_kwargs"].get("blayout") == "ColMajor": + _view_strides = [_rows * _cols, _rows * _cols, _rows * _cols, 1, _rows] + _tile_kwargs = dict(_spec["tile_kwargs"]) + _kernel_name = _spec["kernel_name"] + _case_name = _spec["case_name"] + + def _make(rows=_rows, cols=_cols, view_strides=_view_strides, tile_kwargs=_tile_kwargs, kernel_name=_kernel_name): + @pto.jit( + name=kernel_name, + target="a5", + ) + def _kernel( + src_ptr: pto.ptr(pto.f32, "gm"), + dst_ptr: pto.ptr(pto.f32, "gm"), + ): + _roundtrip_body( + src_ptr, + dst_ptr, + rows=rows, + cols=cols, + view_strides=view_strides, + tile_kwargs=tile_kwargs, + ) + + return _kernel + + _tload_store_kernels[_case_name] = _make() + + +def _make_input(name, shape): + import zlib + + np.random.seed(zlib.crc32(name.encode("utf-8")) & 0xFFFFFFFF) + return np.random.randint(1, 32, size=shape).astype(np.float32) + + +def _make_expected(src): + return np.asarray(src, dtype=np.float32).copy() + + +CASES = [] +for _spec in CASE_SPECS: + _case_name = _spec["case_name"] + _shape = _spec["shape"] + CASES.append( + golden_output_case( + "tload_store_" + _case_name, + _tload_store_kernels[_case_name], + inputs=lambda _case_name=_case_name, _shape=_shape: [_make_input(_case_name, _shape)], + expected=_make_expected, + rtol=1e-6, + atol=1e-6, + ) + ) + + +auto_main(globals()) diff --git a/test/tilelib-st/a5/tmatmul/case.py b/test/tilelib-st/a5/tmatmul/case.py new file mode 100644 index 0000000000..e9d10bf73e --- /dev/null +++ b/test/tilelib-st/a5/tmatmul/case.py @@ -0,0 +1,203 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +# Minimal PTODSL cube/tmatmul pilot for A5. +# Goal: validate plain cube tile.matmul lowering/runtime first, without mixing +# MX-specific scale/bias handling or @pto.cube helper boundaries. + +from pathlib import Path +import sys + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +M = 16 +K = 32 +N = 64 +ELEM_BYTES = 4 + +L1_A_ADDR = 0 +L1_B_ADDR = 4096 +L0A_ADDR = 0 +L0B_ADDR = 0 +L0C_ADDR = 0 + +# This case keeps explicit L1/L0 addresses because the current GM->L1 fractal +# path passes raw MAT pointers into mte_gm_l1_frac. Vector tile cases in this +# directory use automatic tile address allocation. + + +@pto.cube +def cube_matmul_tile( + a_mat: pto.Tile, + b_mat: pto.Tile, + o_tile: pto.Tile, + a_l0a: pto.Tile, + b_l0b: pto.Tile, + c_acc: pto.Tile, +): + m = a_mat.valid_shape[0] + k = a_mat.valid_shape[1] + n = b_mat.valid_shape[1] + + pto.mte_l1_l0a(a_mat.as_ptr(), a_l0a.as_ptr(), m, k) + pto.mte_l1_l0b(b_mat.as_ptr(), b_l0b.as_ptr(), k, n, transpose=True) + pto.set_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=1) + pto.wait_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=1) + pto.tile.matmul(a_l0a, b_l0b, c_acc) + pto.set_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=2) + pto.wait_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=2) + pto.mte_l0c_ub( + c_acc.as_ptr(), + o_tile.as_ptr(), + m, + n, + n, + n, + ) + + +@pto.jit( + name="tmatmul_f32_16x32x64", + kernel_kind="cube", + target="a5", + mode="explicit", + insert_sync=False, +) +def _tmatmul_kernel( + a_ptr: pto.ptr(pto.f32, "gm"), + b_ptr: pto.ptr(pto.f32, "gm"), + c_ptr: pto.ptr(pto.f32, "gm"), +): + a_mat = pto.alloc_tile( + shape=[M, K], + dtype=pto.f32, + memory_space=pto.MemorySpace.MAT, + addr=L1_A_ADDR, + valid_shape=[M, K], + blayout="ColMajor", + slayout="RowMajor", + ) + b_mat = pto.alloc_tile( + shape=[K, N], + dtype=pto.f32, + memory_space=pto.MemorySpace.MAT, + addr=L1_B_ADDR, + valid_shape=[K, N], + blayout="ColMajor", + slayout="RowMajor", + ) + a_l0a = pto.alloc_tile( + shape=[M, K], + dtype=pto.f32, + memory_space=pto.MemorySpace.LEFT, + addr=L0A_ADDR, + valid_shape=[M, K], + blayout="ColMajor", + slayout="RowMajor", + ) + b_l0b = pto.alloc_tile( + shape=[K, N], + dtype=pto.f32, + memory_space=pto.MemorySpace.RIGHT, + addr=L0B_ADDR, + valid_shape=[K, N], + blayout="RowMajor", + slayout="ColMajor", + ) + c_acc = pto.alloc_tile( + shape=[M, N], + dtype=pto.f32, + memory_space=pto.MemorySpace.ACC, + addr=L0C_ADDR, + valid_shape=[M, N], + blayout="ColMajor", + slayout="RowMajor", + fractal_size=1024, + ) + + a_l1_ptr = pto.castptr(pto.ui64(L1_A_ADDR), pto.ptr(pto.f32, "mat")) + b_l1_ptr = pto.castptr(pto.ui64(L1_B_ADDR), pto.ptr(pto.f32, "mat")) + + pto.mte_gm_l1_frac( + a_ptr, + a_l1_ptr, + pto.FractalMode.ND2NZ, + shape=(M, K), + src_layout=(K * ELEM_BYTES,), + dst_group=(1, 1, M, 0), + ctrl=(0, False), + ) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=0) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=0) + pto.mte_l1_l0a(a_l1_ptr, a_l0a.as_ptr(), M, K) + + pto.mte_gm_l1_frac( + b_ptr, + b_l1_ptr, + pto.FractalMode.ND2NZ, + shape=(K, N), + src_layout=(N * ELEM_BYTES,), + dst_group=(1, 1, K, 0), + ctrl=(0, False), + ) + pto.set_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=1) + pto.wait_flag(pto.Pipe.MTE2, pto.Pipe.MTE1, event_id=1) + pto.mte_l1_l0b(b_l1_ptr, b_l0b.as_ptr(), K, N, transpose=True) + + pto.set_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) + pto.wait_flag(pto.Pipe.MTE1, pto.Pipe.M, event_id=0) + pto.tile.matmul(a_l0a, b_l0b, c_acc) + + pto.set_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) + pto.wait_flag(pto.Pipe.M, pto.Pipe.FIX, event_id=1) + pto.mte_l0c_gm( + c_acc.as_ptr(), + c_ptr, + M, + N, + M, + N, + 0, + 0, + layout="nz2nd", + ) + pto.pipe_barrier(pto.Pipe.ALL) + + +def _make_inputs(): + rng = np.random.default_rng(0x7A7A7A71) + a = rng.uniform(-2.0, 2.0, size=(M, K)).astype(np.float32) + b = rng.uniform(-2.0, 2.0, size=(K, N)).astype(np.float32) + return [a, b] + + +def _make_expected(a, b): + return (a @ b).astype(np.float32) + + +CASES = [ + golden_output_case( + "tmatmul_f32_16x32x64", + _tmatmul_kernel, + inputs=_make_inputs, + expected=_make_expected, + rtol=1e-4, + atol=1e-4, + ), +] + + +auto_main(globals()) diff --git a/test/tilelib-st/common.py b/test/tilelib-st/common.py new file mode 100644 index 0000000000..44d455469e --- /dev/null +++ b/test/tilelib-st/common.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from __future__ import annotations + +import argparse +import importlib.util +from pathlib import Path +import sys +import time + +import numpy as np + +_DEVICE = "npu:0" + + +def init_runtime(): + import torch + import torch_npu # noqa: F401 + + torch.npu.config.allow_internal_format = False + torch_npu.npu.set_compile_mode(jit_compile=False) + torch.npu.set_device(_DEVICE) + return torch + + +def npu_stream(torch): + return torch.npu.current_stream()._as_parameter_ # noqa: SLF001 + + +def emit_mlir(*kernels): + from ptodsl import pto + + return pto.merge_jit_modules(*kernels) + + +def _collect_case_kernels(cases: list[dict]) -> list: + kernels = [] + seen = set() + for case in cases: + kernel = case["kernel"] + key = id(kernel) + if key in seen: + continue + seen.add(key) + kernels.append(kernel) + return kernels + + +def _to_numpy_array(value): + return np.array(value, copy=True) + + +def golden_output_case( + name: str, + kernel, + *, + inputs, + expected, + output_shape=None, + output_dtype=None, + output_index: int = -1, + launch_args=None, + rtol: float = 1e-5, + atol: float = 1e-5, +): + """Build a standard single-output TileLib ST case from host inputs and a golden.""" + + def materialize_inputs(): + values = inputs() if callable(inputs) else inputs + return [_to_numpy_array(value) for value in values] + + def materialize_expected(host_inputs): + value = expected(*host_inputs) if callable(expected) else expected + return _to_numpy_array(value) + + def make_case(): + host_inputs = materialize_inputs() + golden = materialize_expected(host_inputs) + out = np.zeros( + output_shape or golden.shape, + dtype=output_dtype or golden.dtype, + ) + if launch_args is None: + return [*host_inputs, out], golden + extra_launch_args = launch_args(*host_inputs) if callable(launch_args) else list(launch_args) + return [*host_inputs, out], golden, extra_launch_args + + def check_case(device_inputs, golden): + actual = device_inputs[output_index].cpu().numpy() + assert_close(actual, golden, rtol=rtol, atol=atol) + + return { + "name": name, + "kernel": kernel, + "make_case": make_case, + "check": check_case, + } + + +def _list_cases(cases: list[dict]) -> None: + for case in cases: + print(case["name"]) + + +def _load_module_from_path(path: Path): + rel = path.resolve().relative_to(Path(__file__).resolve().parent) + module_name = "_tilelib_st_" + "_".join(part.replace("-", "_") for part in rel.with_suffix("").parts) + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load TileLib ST module from {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def discover_case_modules(root: Path | None = None) -> list: + case_root = Path(root) if root is not None else Path(__file__).resolve().parent + modules = [] + paths = [] + seen_paths = set() + for path in sorted(case_root.glob("*.py")): + if path.name not in {"common.py", "__main__.py", "run_tilelib_st.py"} and not path.name.startswith("_"): + resolved = path.resolve() + paths.append(path) + seen_paths.add(resolved) + for path in sorted(case_root.rglob("case.py")): + resolved = path.resolve() + if resolved in seen_paths: + continue + paths.append(path) + seen_paths.add(resolved) + + for path in paths: + module = _load_module_from_path(path) + if getattr(module, "CASES", None): + modules.append(module) + return modules + + +def discover_cases(root: Path | None = None) -> list[dict]: + discovered = [] + seen_names = {} + for module in discover_case_modules(root): + module_path = Path(getattr(module, "__file__", "")) + for case in module.CASES: + name = case["name"] + previous = seen_names.get(name) + if previous is not None: + raise RuntimeError( + f"Duplicate TileLib ST case name {name!r} discovered in {module_path} and {previous}" + ) + seen_names[name] = module_path + discovered.append(case) + return discovered + + +def run_cases(cases: list[dict], *, emit_mlir_fn=None, argv=None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--list", action="store_true", help="list discovered case names and exit") + parser.add_argument("--emit-mlir", action="store_true", help="print merged MLIR module and exit") + args = parser.parse_args(argv) + + if args.list: + _list_cases(cases) + return 0 + + if args.emit_mlir: + if emit_mlir_fn is None: + raise RuntimeError("emit_mlir_fn is required when --emit-mlir is supported") + print(emit_mlir_fn()) + return 0 + + torch = init_runtime() + for case in cases: + name = case["name"] + kernel = case["kernel"] + made_case = case["make_case"]() + if len(made_case) == 2: + inputs, expected = made_case + launch_args = [] + elif len(made_case) == 3: + inputs, expected, launch_args = made_case + else: + raise RuntimeError( + f"TileLib ST case {name!r} make_case() must return 2 or 3 values, got {len(made_case)}" + ) + + device_inputs = [torch.from_numpy(array).to(_DEVICE) for array in inputs] + stream = npu_stream(torch) + + t0 = time.perf_counter() + compiled = kernel.compile() + compile_s = time.perf_counter() - t0 + + t0 = time.perf_counter() + compiled[1, stream](*device_inputs, *launch_args) + torch.npu.synchronize() + launch_s = time.perf_counter() - t0 + + case["check"](device_inputs, expected) + print(f"PASS {name} compile={compile_s:.3f}s launch={launch_s:.3f}s") + + print("All cases passed.") + return 0 + + +def run_module_cases(module_globals: dict, argv=None) -> int: + cases = module_globals.get("CASES") + if not cases: + raise RuntimeError("TileLib ST module must define a non-empty CASES list") + + emit_mlir_fn = module_globals.get("EMIT_MLIR_FN") + if emit_mlir_fn is None: + kernels = module_globals.get("KERNELS") + if kernels is None: + kernels = _collect_case_kernels(cases) + emit_mlir_fn = lambda: emit_mlir(*kernels) + + return run_cases(cases, emit_mlir_fn=emit_mlir_fn, argv=argv) + + +def run_discovered_cases(root: Path | None = None, argv=None) -> int: + cases = discover_cases(root) + if not cases: + raise RuntimeError("No TileLib ST cases discovered") + return run_cases( + cases, + emit_mlir_fn=lambda: emit_mlir(*_collect_case_kernels(cases)), + argv=argv, + ) + + +def auto_main(module_globals: dict, argv=None) -> None: + if module_globals.get("__name__") != "__main__": + return + raise SystemExit(run_module_cases(module_globals, argv=argv)) + + +def assert_close(actual, expected, *, rtol=1e-5, atol=1e-5): + np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol) diff --git a/test/tilelib-st/run_tilelib_st.py b/test/tilelib-st/run_tilelib_st.py new file mode 100644 index 0000000000..e425549c22 --- /dev/null +++ b/test/tilelib-st/run_tilelib_st.py @@ -0,0 +1,36 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. + +from pathlib import Path +import sys + + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parent)) + +from common import run_discovered_cases + + +def _resolve_case_root(arg: str | None) -> Path: + if arg is None: + return Path(__file__).resolve().parent + root = Path(arg) + if root.is_absolute(): + return root + repo_root = Path(__file__).resolve().parents[2] + return (repo_root / root).resolve() + + +def main() -> int: + root = _resolve_case_root(sys.argv[1] if len(sys.argv) > 1 else None) + return run_discovered_cases(root, argv=sys.argv[2:]) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/ptoas/driver.cpp b/tools/ptoas/driver.cpp index 6cd1add149..658ba24aef 100644 --- a/tools/ptoas/driver.cpp +++ b/tools/ptoas/driver.cpp @@ -305,6 +305,12 @@ static bool isBackendPartitionedContainer(ModuleOp module) { [](Operation &op) { return isa(op); }); } +static bool isUserVisibleIROutputRequested() { + return mlir::pto::emitMlirIR || mlir::pto::emitVPTO || + mlir::pto::emitVPTOLLVMDialect || mlir::pto::ptoPrintSeamIR || + !mlir::pto::ptoSeamIRFile.empty(); +} + static SmallVector collectImportedPeerNames(ModuleOp module) { SmallVector names; module.walk([&](pto::ImportReservedBufferOp importOp) { @@ -962,12 +968,35 @@ LogicalResult EmitCBackendJob::run(PTOASContext &context) { } LogicalResult VPTOBackendJob::run(PTOASContext &context) { + OwningOpRef singleChildJobModule; + OwningOpRef *compileUnit = &module; ModuleOp op = module.get(); op->setAttr("pto.backend", StringAttr::get(op.getContext(), "vpto")); + SmallVector children(op.getOps()); + // PTODSL emits a backend-partitioned outer container even when there is only + // one child module. For object compilation, the actual VPTO compile unit is + // the normalized child job module with outer attributes/imports folded in. + // Keep user-visible IR output modes on the original container so dump flags + // still reflect the input module structure the user asked to inspect. + if (!isUserVisibleIROutputRequested() && children.size() == 1 && + isBackendPartitionedContainer(op) && + children.front()->hasAttr(mlir::pto::FunctionKernelKindAttr::name)) { + FailureOr> jobModuleOr = + buildBackendChildCompileUnit(op, children.front()); + if (failed(jobModuleOr)) + return failure(); + singleChildJobModule = std::move(*jobModuleOr); + singleChildJobModule.get()->setAttr( + "pto.backend", + StringAttr::get(singleChildJobModule.get()->getContext(), "vpto")); + compileUnit = &singleChildJobModule; + op = singleChildJobModule.get(); + } + bool emitHostStub = hasPTOEntry(op); if (mlir::pto::compilePTOASModule( - module, context, mlir::pto::PTOBackend::VPTO, result, + *compileUnit, context, mlir::pto::PTOBackend::VPTO, result, emitHostStub) != 0) return failure(); if (result.kind == mlir::pto::PTOASCompileResultKind::Text) @@ -1061,10 +1090,7 @@ static LogicalResult resolveSingleBackend( } SmallVector children(module.getOps()); - bool debugIROutputRequested = - mlir::pto::emitMlirIR || mlir::pto::emitVPTO || mlir::pto::ptoPrintSeamIR || - !mlir::pto::ptoSeamIRFile.empty(); - if (!debugIROutputRequested && children.size() > 1) { + if (!isUserVisibleIROutputRequested() && children.size() > 1) { if (!isBackendPartitionedContainer(module)) { llvm::errs() << "Error: mixed pto.backend fatobj mode expects either a " "single module or an outer module containing only child "