Feature/a2a3 llvm

_ptoas_build_backend.py

@@ -35,18 +35,41 @@
 from pathlib import Path
 
 _REPO = Path(__file__).parent.resolve()
-_LLVM_BUILD_DIR = Path(
-    os.environ.get("LLVM_BUILD_DIR",
-                   "/llvm-workspace/llvm-project/build-shared")
-)
+
+def _find_llvm_dir():
+    """Return an LLVM install or build-tree prefix, resolving in order:
+
+    1. ``LLVM_BUILD_DIR`` / ``LLVM_INSTALL_DIR`` env vars
+    2. Auto-detect common install locations by probing ``lib/cmake/llvm``
+    3. Default build-tree path
+    """
+    for key in ("LLVM_BUILD_DIR", "LLVM_INSTALL_DIR"):
+        if key in os.environ:
+            return Path(os.environ[key])
+
+    for cand in ("/usr/local/llvm", "/usr/local/Ascend/latest/compiler",
+                 "/opt/llvm"):
+        if (Path(cand) / "lib" / "cmake" / "llvm").is_dir():
+            return Path(cand)
+
+    return Path("/llvm-workspace/llvm-project/build-shared")
+
+
+_LLVM_BUILD_DIR = _find_llvm_dir()
 _PTO_INSTALL_DIR = Path(
     os.environ.get("PTO_INSTALL_DIR", str(_REPO / "install"))
 )
 _BUILD_DIR = Path(os.environ.get("PTO_BUILD_DIR", str(_REPO / "build")))
 _PTODSL_SOURCE_ROOT = _REPO / "ptodsl"
-_MLIR_PY_PKG = (
-    _LLVM_BUILD_DIR / "tools" / "mlir" / "python_packages" / "mlir_core"
-)
+_MLIR_PY_PKG = None
+if "MLIR_PYTHON_PACKAGE_DIR" in os.environ:
+    _MLIR_PY_PKG = Path(os.environ["MLIR_PYTHON_PACKAGE_DIR"])
+elif "LLVM_INSTALL_DIR" in os.environ:
+    _MLIR_PY_PKG = Path(os.environ["LLVM_INSTALL_DIR"]) / "python_packages" / "mlir_core"
+else:
+    _installed = _LLVM_BUILD_DIR / "python_packages" / "mlir_core"
+    _build_tree = _LLVM_BUILD_DIR / "tools" / "mlir" / "python_packages" / "mlir_core"
+    _MLIR_PY_PKG = _installed if _installed.exists() else _build_tree
 _WHEEL_DIST_DIR = _BUILD_DIR / "wheel-dist"
 
 
@@ -124,7 +147,7 @@ def _should_use_linux_hardening_cache() -> bool:
     return sys.platform.startswith("linux")
 
 
-def _cmake_configure_and_build():
+def _cmake_configure_and_build(skip_install=False):
     """CMake configure + build + install."""
     _BUILD_DIR.mkdir(parents=True, exist_ok=True)
 
@@ -136,6 +159,7 @@ def _cmake_configure_and_build():
         "cmake", "-GNinja",
         f"-S{_REPO}", f"-B{_BUILD_DIR}",
         "-DCMAKE_BUILD_TYPE=Release",
+        "-DPTO_ENABLE_PYTHON_BINDING=ON",
         f"-DLLVM_DIR={_LLVM_BUILD_DIR}/lib/cmake/llvm",
         f"-DMLIR_DIR={_LLVM_BUILD_DIR}/lib/cmake/mlir",
         f"-DPython3_ROOT_DIR={sys.prefix}",
@@ -161,10 +185,11 @@ def _cmake_configure_and_build():
 
     subprocess.check_call(cmake_cmd)
     subprocess.check_call(["cmake", "--build", str(_BUILD_DIR)])
-    subprocess.check_call(
-        ["cmake", "--build", str(_BUILD_DIR), "--target", "install"]
-    )
-    _assert_installed_ptodsl_payload()
+    if not skip_install:
+        subprocess.check_call(
+            ["cmake", "--build", str(_BUILD_DIR), "--target", "install"]
+        )
+        _assert_installed_ptodsl_payload()
 
 
 def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
@@ -209,7 +234,7 @@ def build_editable(wheel_directory, config_settings=None, metadata_directory=Non
     the installed/runtime build outputs. No files are copied into
     site-packages except the .pth file itself.
     """
-    _cmake_configure_and_build()
+    _cmake_configure_and_build(skip_install=True)
     _assert_editable_ptodsl_source()
 
     version = os.environ.get("PTOAS_PYTHON_PACKAGE_VERSION", "0.1.0")
@@ -221,9 +246,11 @@ def build_editable(wheel_directory, config_settings=None, metadata_directory=Non
         # Prefer the repository PTODSL sources so editable installs pick up
         # local Python edits instead of staged/install-tree copies.
         str(_PTODSL_SOURCE_ROOT),
-        # Installed PTOAS runtime overlay (mlir.dialects.pto, _pto, TileOps).
+        # Handwritten Python sources (pto/dialects/pto.py, etc.).
+        str(_REPO / "python"),
+        # Installed PTOAS runtime overlay (TileOps/resources when present).
         str(_PTO_INSTALL_DIR),
-        # Keep the in-tree extension/staging output last as a fallback.
+        # Generated files (_pto.so, _pto_ops_gen.py) under mlir/ namespace.
         str(_BUILD_DIR / "python"),
     ]
 

docker/Dockerfile

@@ -169,7 +169,7 @@ RUN bisheng $BISHENG_OPTS ./tmatmulk.cpp -o ./matmul_kernel.so
 # add torch_npu to faciliate on-device testing during `docker run` (not used at build stage)
 # for available versions see https://pypi.org/project/torch-npu/#history
 RUN pip install --no-cache-dir torch==2.9.0 --index-url https://download.pytorch.org/whl/cpu \
-    && pip install --no-cache-dir torch-npu==2.9.0 \
+    && pip install --no-cache-dir torch-npu==2.9.0.post2 \
     && pip install --no-cache-dir numpy pyyaml
 
 WORKDIR /sources/test/

docker/Dockerfile.dev

@@ -0,0 +1,84 @@
+# PTOAS development environment with NPU support (CANN 9.0.0)
+#
+# Usage:
+#   docker build -f docker/Dockerfile.dev -t ptoas-dev-cann900 .
+#
+# This image provides the toolchain for building and running PTOAS on NPU.
+# It does NOT contain PTOAS itself -- mount PTOAS source at /workspace and
+# run `pip install -e .` inside the container.
+
+# ---- Base: CANN 9.0.0 (provides bisheng, CANN runtime, Python 3.12) ----
+FROM quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.12
+
+# ---- Build args (matching CI) ----
+ARG LLVM_REPO=https://github.com/vpto-dev/llvm-project.git
+ARG LLVM_REF=feature-vpto-llvm21
+ARG PTO_ISA_REPO=https://gitcode.com/cann/pto-isa.git
+ARG PTO_ISA_COMMIT=662d7f2a916d6bbde3109ce4a16ed5c28f5d900a
+
+# ---- Paths ----
+ENV LLVM_SOURCE_DIR=/opt/llvm-project
+ENV LLVM_BUILD_DIR=/opt/llvm-project/build
+ENV LLVM_INSTALL_DIR=/usr/local/llvm
+ENV ASCEND_HOME_PATH=/usr/local/Ascend/cann-9.0.0
+ENV PTO_LIB_PATH=/opt/pto-isa
+
+# =========================================================================
+# 1. System dependencies
+# =========================================================================
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    cmake ninja-build git build-essential clang lld \
+    libedit-dev zlib1g-dev libxml2-dev libzstd-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# =========================================================================
+# 2. Python dependencies
+# =========================================================================
+RUN pip install --no-cache-dir 'pybind11<3' numpy pyyaml pytest nanobind
+
+# =========================================================================
+# 3. Build LLVM/MLIR (vpto fork, feature-vpto branch)
+# =========================================================================
+RUN git clone --depth 1 --branch "${LLVM_REF}" "${LLVM_REPO}" "${LLVM_SOURCE_DIR}"
+
+WORKDIR ${LLVM_SOURCE_DIR}
+
+RUN cmake -G Ninja -S llvm -B "${LLVM_BUILD_DIR}" \
+    -DLLVM_ENABLE_PROJECTS="mlir;clang" \
+    -DBUILD_SHARED_LIBS=ON \
+    -DLLVM_ENABLE_ASSERTIONS=ON \
+    -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
+    -DPython3_EXECUTABLE=python3 \
+    -DCMAKE_BUILD_TYPE=Release \
+    -DLLVM_TARGETS_TO_BUILD="host" \
+    -DCMAKE_INSTALL_PREFIX="${LLVM_INSTALL_DIR}" \
+    -DLLVM_INSTALL_UTILS=ON
+
+RUN ninja -C "${LLVM_BUILD_DIR}"
+
+RUN cmake --install "${LLVM_BUILD_DIR}"
+
+# =========================================================================
+# 4. Clone pto-isa (headers for bisheng compilation of generated kernels)
+# =========================================================================
+RUN git clone "${PTO_ISA_REPO}" "${PTO_LIB_PATH}" \
+    && cd "${PTO_LIB_PATH}" && git checkout "${PTO_ISA_COMMIT}"
+
+# =========================================================================
+# 5. Torch + torch-npu (for on-device kernel launch)
+# =========================================================================
+RUN pip install --no-cache-dir torch==2.9.0 --index-url https://download.pytorch.org/whl/cpu \
+    && pip install --no-cache-dir torch-npu==2.9.0.post2
+
+# =========================================================================
+# 6. Environment
+# =========================================================================
+ENV PATH="${LLVM_INSTALL_DIR}/bin:${ASCEND_HOME_PATH}/compiler/bin:${PATH}"
+ENV LD_LIBRARY_PATH="${LLVM_INSTALL_DIR}/lib:${ASCEND_HOME_PATH}/lib64:${ASCEND_HOME_PATH}/lib64/plugin/opskernel:${ASCEND_HOME_PATH}/lib64/plugin/nnengine:${ASCEND_HOME_PATH}/opp/built-in/op_impl/ai_core/tbe/op_tiling/lib/linux/aarch64:${LD_LIBRARY_PATH}"
+ENV MLIR_PYTHON_ROOT="${LLVM_INSTALL_DIR}/python_packages/mlir_core"
+
+ENV BISHENG_OPTS="-I${PTO_LIB_PATH}/include/pto -fPIC -shared -O2 -std=c++17 --npu-arch=dav-2201 -DMEMORY_BASE"
+
+WORKDIR /workspace
+
+CMD ["/bin/bash"]