[feat]Add torch triton-runtime headers, demo and README-pytorch-AtomGit

Aexamples/libtorch_triton_ascend/README.md+56-0

@@ -0,0 +1,56 @@
	1	+# libtorch_triton_ascend
	2	+
	3	+## 跑demo
	4	+
	5	+demo.cpp 包含三个示例：
	6	+
	7	+1. Add Kernel (explicit grid) - 使用显式 grid 的向量加法 kernel，展示最基本的 Triton kernel 调用方式
	8	+2. Add Kernel (C++ lambda grid) - 使用 C++ lambda 函数计算 grid，类似 Triton 的 `lambda meta`，完全在 C++ 端解析无需 Python 参与
	9	+3. LayerNorm Kernel - LayerNorm 算子示例，展示更复杂的 kernel 参数传递和计算
	10	+
	11	+运行流程：
	12	+
	13	+1）generate libtriton_runtime.so
	14	+
	15	+```bash
	16	+cd triton_runtime
	17	+bash build.sh
	18	+cd ..
	19	+```
	20	+
	21	+2）run demo
	22	+
	23	+```bash
	24	+bash examples/run_demo.sh
	25	+```
	26	+
	27	+3）期望输出
	28	+
	29	+```text
	30	+...
	31	+=== LayerNorm Kernel ===
	32	+[KernelRegistry] 'layer_norm_kernel': def layer_norm_kernel(output_ptr, input_ptr, weight_ptr, bias_ptr, n_rows: [constexpr], n_cols: [constexpr], stride, eps, BLOCK_SIZE: [constexpr])
	33	+C++ BLOCK_SIZE: 512
	34	+C++ triton out[0,:5]: 0.1781
	35	+-1.6806
	36	+-1.2184
	37	+-0.3596
	38	+-0.9843
	39	+[ npuFloatType{5} ]
	40	+C++ torch ref[0,:5]: 0.1781
	41	+-1.6806
	42	+-1.2184
	43	+-0.3596
	44	+-0.9843
	45	+[ npuFloatType{5} ]
	46	+C++: max diff: 7.15256e-07
	47	+C++: match: YES
	48	+ registered: layer_norm_kernel
	49	+ registered: add_kernel
	50	+```
	51	+
	52	+4）打开环境变量查看详细日志
	53	+
	54	+```bash
	55	+export LOG_TORCH_TRITON_RUNTIME=1
	56	+```

Aexamples/libtorch_triton_ascend/examples/CMakeLists.txt+69-0

@@ -0,0 +1,69 @@
	1	+cmake_minimum_required(VERSION 3.18)
	2	+project(demo LANGUAGES CXX)
	3	+
	4	+# Helper function for demo targets
	5	+function(add_demo target_name)
	6	+ add_executable(${target_name} ${target_name}.cpp)
	7	+
	8	+ target_include_directories(${target_name} PRIVATE
	9	+ ${TRITON_RUNTIME_DIR}/include
	10	+ ${PYTORCH_PYTHON_PACKAGES}/include
	11	+ ${PYTORCH_NPU_PACKAGES}/include
	12	+ ${Python3_INCLUDE_DIRS}
	13	+ ${CANN_PATH}/include
	14	+ ${CANN_PATH}/include/experiment
	15	+ ${CANN_PATH}/include/experiment/msprof
	16	+ ${CANN_PATH}/pkg_inc
	17	+ ${CANN_PATH}/pkg_inc/profiling
	18	+ )
	19	+
	20	+ target_link_directories(${target_name} PRIVATE
	21	+ ${TRITON_RUNTIME_DIR}/build
	22	+ ${PYTORCH_PYTHON_PACKAGES}/lib
	23	+ ${PYTORCH_PYTHON_PACKAGES}/../torch.libs
	24	+ ${CANN_PATH}/lib64
	25	+ ${PYTORCH_NPU_PACKAGES}/lib
	26	+ )
	27	+
	28	+ target_link_libraries(${target_name} PRIVATE
	29	+ triton_runtime
	30	+ torch
	31	+ ${TORCH_LIBRARIES}
	32	+ torch_cpu
	33	+ torch_npu
	34	+ c10
	35	+ dl
	36	+ Python3::Python
	37	+ runtime
	38	+ ascendcl
	39	+ profapi
	40	+ ssl
	41	+ crypto
	42	+ )
	43	+
	44	+ set_target_properties(${target_name} PROPERTIES
	45	+ BUILD_RPATH "${TRITON_RUNTIME_DIR}/build;${PYTORCH_PYTHON_PACKAGES}/lib;${PYTORCH_NPU_PACKAGES}/lib;${CANN_PATH}/lib64"
	46	+ INSTALL_RPATH "${TRITON_RUNTIME_DIR}/build;${PYTORCH_PYTHON_PACKAGES}/lib;${PYTORCH_NPU_PACKAGES}/lib;${CANN_PATH}/lib64"
	47	+ )
	48	+endfunction()
	49	+
	50	+set(CMAKE_CXX_STANDARD 17)
	51	+set(CMAKE_CXX_STANDARD_REQUIRED ON)
	52	+
	53	+# ──── Paths (override with -D flags if needed) ────
	54	+
	55	+set(PYTORCH_PYTHON_PACKAGES "${PYTORCH_PYTHON_PACKAGES}" CACHE PATH "torch Python packages directory")
	56	+set(PYTORCH_NPU_PACKAGES "${PYTORCH_NPU_PACKAGES}" CACHE PATH "torch_npu Python packages directory")
	57	+set(CANN_PATH "${CANN_PATH}" CACHE PATH "CANN installation root")
	58	+
	59	+message("CMAKE_CURRENT_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}")
	60	+set(TRITON_RUNTIME_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../triton_runtime" CACHE PATH "Path to triton_runtime project root")
	61	+
	62	+set(PYTHON_EXECUTABLE "${PYTHON_EXECUTABLE}" CACHE FILEPATH "Python executable path")
	63	+
	64	+find_package(Python3 3.10 COMPONENTS Interpreter Development REQUIRED)
	65	+find_package(Torch REQUIRED)
	66	+
	67	+# ──── Demo executables ────
	68	+
	69	+add_demo(demo)

Aexamples/libtorch_triton_ascend/examples/demo.cpp+158-0

@@ -0,0 +1,158 @@
	1	+#include "triton_runtime.h"
	2	+#include <torch/torch.h>
	3	+#include <iostream>
	4	+#include <torch_npu/csrc/core/npu/NPUStream.h>
	5	+
	6	+using namespace triton_runtime;
	7	+
	8	+void run_add_kernel() {
	9	+ auto& rt = TritonRuntime::instance();
	10	+
	11	+ // register kernel
	12	+ auto s = rt.register_kernel("examples/my_kernels.py", "add_kernel");
	13	+ if (!s.ok()) {
	14	+ std::cerr << "Register failed: " << s.error_message() << std::endl;
	15	+ return;
	16	+ }
	17	+
	18	+ rt.print_kernel_signature("add_kernel");
	19	+
	20	+ // Prepare data on NPU
	21	+ auto x = torch::rand({1000}, torch::kFloat32).to("npu");
	22	+ auto y = torch::rand({1000}, torch::kFloat32).to("npu");
	23	+ auto out = torch::empty_like(x);
	24	+ int n_elements = 1000;
	25	+ int BLOCK_SIZE = 1024;
	26	+
	27	+ auto grid = (n_elements + BLOCK_SIZE - 1) / BLOCK_SIZE;
	28	+
	29	+ // Launch Triton kernel
	30	+ s = rt.run("add_kernel", Grid(grid, 1, 1), x, y, out, n_elements, BLOCK_SIZE);
	31	+ if (!s.ok()) {
	32	+ std::cerr << "Launch failed: " << s.error_message() << std::endl;
	33	+ return;
	34	+ }
	35	+
	36	+ // validate result
	37	+ auto libtorch_result = torch::add(x, y);
	38	+ std::cout << "C++ libtorch add[0:5]: " << libtorch_result.slice(0, 0, 5) << std::endl;
	39	+ std::cout << "C++ triton out[0:5]: " << out.slice(0, 0, 5) << std::endl;
	40	+ auto diff = (libtorch_result - out).abs().max().item<float>();
	41	+ std::cout << "C++: max diff: " << diff << std::endl;
	42	+ std::cout << "C++: match: " << (diff < 1e-5 ? "YES" : "NO") << std::endl;
	43	+}
	44	+
	45	+// C++ lambda grid — like Triton's lambda meta: (...) but resolved entirely
	46	+// in C++ without Python involvement.
	47	+void run_add_kernel_lambda_grid() {
	48	+ auto& rt = TritonRuntime::instance();
	49	+
	50	+ auto s = rt.register_kernel("examples/my_kernels.py", "add_kernel");
	51	+ if (!s.ok()) {
	52	+ std::cerr << "Register failed: " << s.error_message() << std::endl;
	53	+ return;
	54	+ }
	55	+
	56	+ auto x = torch::rand({1000}, torch::kFloat32).to("npu");
	57	+ auto y = torch::rand({1000}, torch::kFloat32).to("npu");
	58	+ auto out = torch::empty_like(x);
	59	+ int n_elements = 1000;
	60	+ int BLOCK_SIZE = 1024;
	61	+
	62	+ // meta["key"] returns variant<int32_t,float,bool>, same as Triton's bound_args.
	63	+ GridFn grid_fn = [](BoundArgs& meta) -> Grid {
	64	+ int n = meta["n_elements"];
	65	+ int bs = meta["BLOCK_SIZE"];
	66	+ return Grid((n + bs - 1) / bs, 1, 1);
	67	+ };
	68	+
	69	+ s = rt.run("add_kernel", Grid(grid_fn), x, y, out,
	70	+ n_elements, BLOCK_SIZE);
	71	+ if (!s.ok()) {
	72	+ std::cerr << "Launch failed: " << s.error_message() << std::endl;
	73	+ return;
	74	+ }
	75	+
	76	+ auto libtorch_result = torch::add(x, y);
	77	+ std::cout << "C++ libtorch add[0:5]: " << libtorch_result.slice(0, 0, 5) << std::endl;
	78	+ std::cout << "C++ triton out[0:5]: " << out.slice(0, 0, 5) << std::endl;
	79	+ auto diff = (libtorch_result - out).abs().max().item<float>();
	80	+ std::cout << "C++: max diff: " << diff << std::endl;
	81	+ std::cout << "C++: match: " << (diff < 1e-5 ? "YES" : "NO") << std::endl;
	82	+}
	83	+
	84	+void run_layer_norm_kernel() {
	85	+ auto& rt = TritonRuntime::instance();
	86	+
	87	+ // register kernel
	88	+ auto s = rt.register_kernel("examples/my_kernels.py", "layer_norm_kernel");
	89	+ if (!s.ok()) {
	90	+ std::cerr << "Register failed: " << s.error_message() << std::endl;
	91	+ return;
	92	+ }
	93	+
	94	+ rt.print_kernel_signature("layer_norm_kernel");
	95	+
	96	+ // Prepare data on NPU
	97	+ int n_rows = 128;
	98	+ int n_cols = 512;
	99	+ float eps = 1e-5;
	100	+ auto x = torch::randn({n_rows, n_cols}, torch::kFloat32).to("npu");
	101	+ auto weight = torch::ones({n_cols}, torch::kFloat32).to("npu");
	102	+ auto bias = torch::zeros({n_cols}, torch::kFloat32).to("npu");
	103	+ auto out = torch::empty_like(x);
	104	+ int64_t stride = x.stride(0);
	105	+
	106	+ // BLOCK_SIZE = next_power_of_2(n_cols)
	107	+ int BLOCK_SIZE = 1;
	108	+ while (BLOCK_SIZE < n_cols) BLOCK_SIZE *= 2;
	109	+
	110	+ std::cout << "C++ BLOCK_SIZE: " << BLOCK_SIZE << std::endl;
	111	+
	112	+ // Launch Triton kernel
	113	+ s = rt.run("layer_norm_kernel", Grid(n_rows, 1, 1),
	114	+ out, x, weight, bias,
	115	+ n_rows, n_cols, stride, eps,
	116	+ BLOCK_SIZE);
	117	+ if (!s.ok()) {
	118	+ std::cerr << "Launch failed: " << s.error_message() << std::endl;
	119	+ return;
	120	+ }
	121	+
	122	+ // validate result
	123	+ auto ref = torch::layer_norm(x, {n_cols}, weight, bias, eps);
	124	+ std::cout << "C++ triton out[0,:5]: " << out.index({0, torch::indexing::Slice(0, 5)}) << std::endl;
	125	+ std::cout << "C++ torch ref[0,:5]: " << ref.index({0, torch::indexing::Slice(0, 5)}) << std::endl;
	126	+ auto diff = (ref - out).abs().max().item<float>();
	127	+ std::cout << "C++: max diff: " << diff << std::endl;
	128	+ std::cout << "C++: match: " << (diff < 1e-3 ? "YES" : "NO") << std::endl;
	129	+}
	130	+
	131	+int main() {
	132	+ // List registered kernels
	133	+ auto& rt = TritonRuntime::instance();
	134	+
	135	+ // simple case — explicit grid
	136	+ std::cout << "=== Add Kernel (explicit grid) ===" << std::endl;
	137	+ run_add_kernel();
	138	+
	139	+
	140	+ // lambda grid — C++ callable, resolves without Python
	141	+ std::cout << "\n=== Add Kernel (C++ lambda grid) ===" << std::endl;
	142	+ run_add_kernel_lambda_grid();
	143	+
	144	+ for (const auto& name : rt.list_kernels()) {
	145	+ std::cout << " registered: " << name << std::endl;
	146	+ }
	147	+
	148	+
	149	+ // complicate case
	150	+ std::cout << "\n=== LayerNorm Kernel ===" << std::endl;
	151	+ run_layer_norm_kernel();
	152	+
	153	+ for (const auto& name : rt.list_kernels()) {
	154	+ std::cout << " registered: " << name << std::endl;
	155	+ }
	156	+
	157	+ return 0;
	158	+}

Aexamples/libtorch_triton_ascend/examples/my_kernels.py+63-0

@@ -0,0 +1,63 @@
	1	+import triton
	2	+import triton.language as tl
	3	+
	4	+@triton.jit
	5	+def add_kernel(x_ptr: tl.tensor,
	6	+ y_ptr: tl.tensor,
	7	+ output_ptr,
	8	+ n_elements,
	9	+ BLOCK_SIZE: tl.constexpr,
	10	+ ):
	11	+ """向量加法 kernel（参考 test_libtorch_ffmh）"""
	12	+ pid = tl.program_id(axis=0)
	13	+ block_start = pid * BLOCK_SIZE
	14	+ offsets = block_start + tl.arange(0, BLOCK_SIZE)
	15	+ mask = offsets < n_elements
	16	+ x = tl.load(x_ptr + offsets, mask=mask)
	17	+ y = tl.load(y_ptr + offsets, mask=mask)
	18	+ output = x + y
	19	+ tl.store(output_ptr + offsets, output, mask=mask)
	20	+
	21	+
	22	+@triton.jit
	23	+def layer_norm_kernel(
	24	+ output_ptr,
	25	+ input_ptr,
	26	+ weight_ptr,
	27	+ bias_ptr,
	28	+ n_rows: tl.constexpr,
	29	+ n_cols: tl.constexpr,
	30	+ stride,
	31	+ eps,
	32	+ BLOCK_SIZE: tl.constexpr,
	33	+):
	34	+ """Fused LayerNorm kernel: output = (x - mean) / sqrt(var + eps) * weight + bias
	35	+
	36	+ 每个 program 处理一行数据，计算该行的均值和方差后归一化。
	37	+ 支持 2D 输入 (n_rows x n_cols)，weight/bias 为长度 n_cols 的向量。
	38	+ """
	39	+ row_idx = tl.program_id(0)
	40	+ row_start = row_idx * stride
	41	+ col_offsets = tl.arange(0, BLOCK_SIZE)
	42	+ mask = col_offsets < n_cols
	43	+
	44	+ # 加载一行数据
	45	+ x = tl.load(input_ptr + row_start + col_offsets, mask=mask, other=0.0)
	46	+
	47	+ # 计算均值
	48	+ x_mean = tl.sum(x, axis=0) / n_cols
	49	+
	50	+ # 计算方差
	51	+ x_centered = x - x_mean
	52	+ x_var = tl.sum(x_centered * x_centered, axis=0) / n_cols
	53	+
	54	+ # 归一化
	55	+ rstd = 1.0 / tl.sqrt(x_var + eps)
	56	+ x_norm = x_centered * rstd
	57	+
	58	+ # 仿射变换: weight * x_norm + bias
	59	+ weight = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0)
	60	+ bias = tl.load(bias_ptr + col_offsets, mask=mask, other=0.0)
	61	+ output = weight * x_norm + bias
	62	+
	63	+ tl.store(output_ptr + row_start + col_offsets, output, mask=mask)

Aexamples/libtorch_triton_ascend/examples/run_demo.sh+58-0

@@ -0,0 +1,58 @@
	1	+#!/bin/bash
	2	+set -e
	3	+
	4	+export TORCH_DEVICE_BACKEND_AUTOLOAD=0
	5	+
	6	+SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
	7	+BUILD_DIR=${SCRIPT_DIR}/build
	8	+DEMO_BIN=${BUILD_DIR}/demo
	9	+
	10	+# ── Environment ──
	11	+
	12	+PYTHON3="${PYTHON3:-python3}"
	13	+PYTHONHOME=$($PYTHON3 -c 'import sys; print(sys.prefix)')
	14	+TORCH_PATH=$($PYTHON3 -c 'import torch; print(torch.__path__[0])')
	15	+TORCH_NPU_PATH=$($PYTHON3 -c 'import torch_npu; print(torch_npu.__path__[0])')
	16	+CANN_ROOT="${ASCEND_HOME_PATH:-${ASCEND_HOME:-}}"
	17	+
	18	+if [ -z "$CANN_ROOT" ]; then
	19	+ echo "[ERROR] Set ASCEND_HOME_PATH or ASCEND_HOME first"
	20	+ exit 1
	21	+fi
	22	+
	23	+TORCH_CMAKE_PREFIX=$($PYTHON3 -c 'import torch; print(torch.utils.cmake_prefix_path)' 2>/dev/null)
	24	+PYBIND11_CMAKE_DIR=$($PYTHON3 -m pybind11 --cmakedir 2>/dev/null)
	25	+
	26	+# ── Build ──
	27	+
	28	+echo "=== Build demo ==="
	29	+echo "PYTHON3: $(which $PYTHON3)"
	30	+echo "PYTHONHOME: ${PYTHONHOME}"
	31	+echo "TORCH_CMAKE_PREFIX: ${TORCH_CMAKE_PREFIX}"
	32	+echo "TORCH_PATH: ${TORCH_PATH}"
	33	+echo "TORCH_NPU_PATH: ${TORCH_NPU_PATH}"
	34	+cmake -S "${SCRIPT_DIR}" -B "${BUILD_DIR}" -DCMAKE_BUILD_TYPE=Release \
	35	+ -DPYTORCH_PYTHON_PACKAGES="${TORCH_PATH}" \
	36	+ -DPYTORCH_NPU_PACKAGES="${TORCH_NPU_PATH}" \
	37	+ -DCANN_PATH="${CANN_ROOT}" \
	38	+ -DPYTHON_EXECUTABLE="$(which $PYTHON3)" \
	39	+ -DCMAKE_PREFIX_PATH="${TORCH_CMAKE_PREFIX};${PYBIND11_CMAKE_DIR}"
	40	+cmake --build "${BUILD_DIR}" -j$(nproc)
	41	+
	42	+# ── Run ──
	43	+
	44	+echo ""
	45	+echo "=== Run demo ==="
	46	+echo ""
	47	+
	48	+TORCH_LIB="${TORCH_PATH}/lib"
	49	+TORCH_NPU_LIB="${TORCH_NPU_PATH}/lib"
	50	+CANN_LIB="${CANN_ROOT}/lib64"
	51	+RT_BUILD="${SCRIPT_DIR}/../build"
	52	+
	53	+PROJECT_ROOT="${SCRIPT_DIR}/../.."
	54	+
	55	+PYTHONHOME="${PYTHONHOME}" \
	56	+PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}" \
	57	+LD_LIBRARY_PATH="${RT_BUILD}:${TORCH_LIB}:${TORCH_NPU_LIB}:${CANN_LIB}:${LD_LIBRARY_PATH}" \
	58	+ "$DEMO_BIN"

Aexamples/libtorch_triton_ascend/triton_runtime/CMakeLists.txt+91-0

@@ -0,0 +1,91 @@
	1	+cmake_minimum_required(VERSION 3.18)
	2	+project(triton_runtime VERSION 0.1.0 LANGUAGES CXX)
	3	+
	4	+set(CMAKE_CXX_STANDARD 17)
	5	+set(CMAKE_CXX_STANDARD_REQUIRED ON)
	6	+
	7	+set(PYTORCH_PYTHON_PACKAGES "${PYTORCH_PYTHON_PACKAGES}" CACHE PATH "torch Python packages directory")
	8	+set(PYTORCH_NPU_PACKAGES "${PYTORCH_NPU_PACKAGES}" CACHE PATH "torch_npu Python packages directory")
	9	+set(CANN_PATH "${CANN_PATH}" CACHE PATH "CANN installation root")
	10	+
	11	+set(PYTHON_EXECUTABLE "${PYTHON_EXECUTABLE}" CACHE FILEPATH "Python executable path")
	12	+
	13	+find_package(Python3 3.10 COMPONENTS Interpreter Development REQUIRED)
	14	+find_package(Torch REQUIRED)
	15	+find_package(pybind11 REQUIRED)
	16	+
	17	+# ──── Library ────
	18	+
	19	+set(LIB_SOURCES
	20	+ src/arg_info.cpp
	21	+ src/python_env.cpp
	22	+ src/kernel_registry.cpp
	23	+ src/kernel_compiler.cpp
	24	+ src/kernel_cache.cpp
	25	+ src/triton_runtime.cpp
	26	+)
	27	+
	28	+add_library(triton_runtime SHARED ${LIB_SOURCES})
	29	+
	30	+target_include_directories(triton_runtime
	31	+ PUBLIC
	32	+ ${CMAKE_CURRENT_SOURCE_DIR}/include
	33	+ PRIVATE
	34	+ ${PYTORCH_PYTHON_PACKAGES}/include
	35	+ ${PYTORCH_NPU_PACKAGES}/include
	36	+ ${Python3_INCLUDE_DIRS}
	37	+ ${CANN_PATH}/include
	38	+ ${CANN_PATH}/include/experiment
	39	+ ${CANN_PATH}/include/experiment/msprof
	40	+ ${CANN_PATH}/pkg_inc
	41	+ ${CANN_PATH}/pkg_inc/profiling
	42	+)
	43	+
	44	+target_link_directories(triton_runtime PRIVATE
	45	+ ${PYTORCH_PYTHON_PACKAGES}/lib
	46	+ ${PYTORCH_PYTHON_PACKAGES}/../torch.libs
	47	+ ${CANN_PATH}/lib64
	48	+ ${PYTORCH_NPU_PACKAGES}/lib
	49	+)
	50	+
	51	+target_link_libraries(triton_runtime
	52	+ PUBLIC
	53	+ torch
	54	+ ${TORCH_LIBRARIES}
	55	+ PRIVATE
	56	+ pybind11::embed
	57	+ torch_cpu
	58	+ torch_npu
	59	+ c10
	60	+ dl
	61	+ Python3::Python
	62	+ runtime
	63	+ ascendcl
	64	+ profapi
	65	+ ssl
	66	+ crypto
	67	+)
	68	+
	69	+target_compile_options(triton_runtime PRIVATE -Wall -Wextra -fPIC)
	70	+set_target_properties(triton_runtime PROPERTIES
	71	+ OUTPUT_NAME triton_runtime
	72	+ VERSION ${PROJECT_VERSION}
	73	+ SOVERSION 0
	74	+)
	75	+
	76	+# ──── Install ────
	77	+
	78	+install(TARGETS triton_runtime
	79	+ LIBRARY DESTINATION lib
	80	+)
	81	+install(DIRECTORY include/
	82	+ DESTINATION include
	83	+)
	84	+
	85	+# ──── Tests ────
	86	+
	87	+option(TRITON_RUNTIME_BUILD_TESTS "Build tests" OFF)
	88	+if(TRITON_RUNTIME_BUILD_TESTS)
	89	+ enable_testing()
	90	+ add_subdirectory(test)
	91	+endif()

Aexamples/libtorch_triton_ascend/triton_runtime/build.sh+57-0

@@ -0,0 +1,57 @@
	1	+#!/bin/bash
	2	+
	3	+SCRIPT_DIR=$(cd "$(dirname "$0")" && pwd)
	4	+BUILD_DIR=${SCRIPT_DIR}/build
	5	+
	6	+# ── Environment ──
	7	+if [ -z "$ASCEND_HOME_PATH" ] && [ -z "$ASCEND_HOME" ]; then
	8	+ echo "[ERROR] Set ASCEND_HOME_PATH or ASCEND_HOME first"
	9	+ exit 1
	10	+fi
	11	+CANN_ROOT="${ASCEND_HOME_PATH:-$ASCEND_HOME}"
	12	+
	13	+PYTHON3="${PYTHON3:-python3}"
	14	+if ! command -v $PYTHON3 &>/dev/null; then
	15	+ echo "[ERROR] python3 not found"
	16	+ exit 1
	17	+fi
	18	+
	19	+TORCH_PATH=$($PYTHON3 -c 'import torch; print(torch.__path__[0])')
	20	+TORCH_NPU_PATH=$($PYTHON3 -c 'import torch_npu; print(torch_npu.__path__[0])')
	21	+
	22	+TORCH_CMAKE_PREFIX=$($PYTHON3 -c 'import torch; print(torch.utils.cmake_prefix_path)' 2>/dev/null)
	23	+if [ -z "$TORCH_CMAKE_PREFIX" ]; then
	24	+ echo "[ERROR] torch not installed or not importable"
	25	+ exit 1
	26	+fi
	27	+
	28	+PYBIND11_CMAKE_DIR=$($PYTHON3 -m pybind11 --cmakedir 2>/dev/null)
	29	+if [ -z "$PYBIND11_CMAKE_DIR" ]; then
	30	+ echo "[ERROR] pybind11 not installed or not importable"
	31	+ exit 1
	32	+fi
	33	+
	34	+echo "=== Build triton_runtime ==="
	35	+echo " CANN_ROOT: $CANN_ROOT"
	36	+echo " Python3: $PYTHON3"
	37	+echo " Torch cmake prefix: $TORCH_CMAKE_PREFIX"
	38	+echo " pybind11 cmake dir: $PYBIND11_CMAKE_DIR"
	39	+
	40	+# ── CMake Configure ──
	41	+mkdir -p "$BUILD_DIR"
	42	+cd "$BUILD_DIR"
	43	+
	44	+cmake .. \
	45	+ -DCMAKE_BUILD_TYPE=Debug \
	46	+ -DPYTHON_EXECUTABLE="$(which $PYTHON3)" \
	47	+ -DPYTORCH_PYTHON_PACKAGES="${TORCH_PATH}" \
	48	+ -DPYTORCH_NPU_PACKAGES="${TORCH_NPU_PATH}" \
	49	+ -DCANN_PATH="${CANN_ROOT}" \
	50	+ -DCMAKE_PREFIX_PATH="$TORCH_CMAKE_PREFIX;$PYBIND11_CMAKE_DIR"
	51	+
	52	+# ── Build ──
	53	+make -j$(nproc)
	54	+
	55	+echo ""
	56	+echo "=== Build succeeded ==="
	57	+echo " SO: $BUILD_DIR/libtriton_runtime.so"

Aexamples/libtorch_triton_ascend/triton_runtime/include/arg_info.h+36-0

@@ -0,0 +1,36 @@
	1	+#pragma once
	2	+#include <torch/torch.h>
	3	+#include <string>
	4	+#include <variant>
	5	+
	6	+namespace triton_runtime {
	7	+
	8	+class ArgInfo {
	9	+public:
	10	+ using ScalarVariant = std::variant<int32_t, float, bool>;
	11	+
	12	+ ArgInfo(const at::Tensor& tensor);
	13	+ ArgInfo(int32_t val);
	14	+ ArgInfo(int64_t val);
	15	+ ArgInfo(float val);
	16	+ ArgInfo(double val);
	17	+ ArgInfo(bool val);
	18	+
	19	+ bool is_pointer() const;
	20	+ std::string scalar_value() const;
	21	+
	22	+ void* device_ptr() const;
	23	+ const at::Tensor& tensor() const;
	24	+ const ScalarVariant& scalar() const;
	25	+ std::vector<int64_t> shape() const;
	26	+ size_t scalar_size() const;
	27	+
	28	+private:
	29	+ void* ptr_ = nullptr;
	30	+ ScalarVariant scalar_ = int32_t(0);
	31	+ std::vector<int64_t> shape_;
	32	+ at::Tensor tensor_;
	33	+ size_t scalar_size_ = sizeof(int32_t);
	34	+};
	35	+
	36	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/grid.h+50-0

@@ -0,0 +1,50 @@
	1	+#pragma once
	2	+
	3	+#include <functional>
	4	+#include <string>
	5	+#include <unordered_map>
	6	+#include <variant>
	7	+
	8	+namespace triton_runtime {
	9	+
	10	+// Wraps ArgInfo::ScalarVariant with implicit conversion, so the lambda can do
	11	+// int n = meta["n_elements"]; // operator int()
	12	+struct ScalarValue {
	13	+ using Var = std::variant<int32_t, float, bool>;
	14	+ Var v;
	15	+
	16	+ ScalarValue() = default;
	17	+ ScalarValue(const Var& v) : v(v) {}
	18	+ ScalarValue(int32_t x) : v(x) {}
	19	+ ScalarValue(float x) : v(x) {}
	20	+ ScalarValue(bool x) : v(x) {}
	21	+
	22	+ operator int() const { return std::visit([](auto x) { return int(x); }, v); }
	23	+ operator float() const { return std::visit([](auto x) { return float(x); }, v); }
	24	+ operator bool() const { return std::visit([](auto x) { return bool(x); }, v); }
	25	+};
	26	+
	27	+using BoundArgs = std::unordered_map<std::string, ScalarValue>;
	28	+
	29	+struct Grid;
	30	+using GridFn = std::function<Grid(BoundArgs&)>;
	31	+
	32	+// Grid is either a fixed (x,y,z) or a callable that resolves (x,y,z)
	33	+// from bound_args. Both modes live in the same type — no extra wrapper.
	34	+struct Grid {
	35	+ int x = 1, y = 1, z = 1;
	36	+ GridFn fn;
	37	+
	38	+ Grid() = default;
	39	+ Grid(int x, int y = 1, int z = 1) : x(x), y(y), z(z) {}
	40	+ Grid(GridFn f) : fn(std::move(f)) {}
	41	+
	42	+ bool is_callable() const { return !!fn; }
	43	+
	44	+ Grid resolve(BoundArgs& args) const {
	45	+ if (fn) return fn(args);
	46	+ return *this;
	47	+ }
	48	+};
	49	+
	50	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/kernel_cache.h+77-0

@@ -0,0 +1,77 @@
	1	+#pragma once
	2	+#include <string>
	3	+#include <vector>
	4	+#include <memory>
	5	+
	6	+#include <unordered_map>
	7	+#include <shared_mutex>
	8	+#include <sstream>
	9	+#include <dlfcn.h>
	10	+#include <cstdint>
	11	+#include "grid.h"
	12	+#include "arg_info.h"
	13	+#include "kernel_registry.h"
	14	+
	15	+namespace triton_runtime {
	16	+
	17	+// dlopen-based launcher: the generated launcher .so exposes this extern "C" entry point
	18	+typedef void (*TritonLaunchKernelFn)(
	19	+ const char* kernelName, const void* func, void* stream,
	20	+ int gridX, int gridY, int gridZ,
	21	+ const int64_t* shapes_data, const int* shape_dims, int num_tensors,
	22	+ const int* tensor_kinds,
	23	+ void* const* kernel_args, const size_t* arg_sizes, int num_args);
	24	+
	25	+struct CompiledKernelEntry {
	26	+ std::string kernel_name;
	27	+ void* kernel_func_ptr = nullptr;
	28	+ std::string launcher_so_path;
	29	+ void* launcher_so_handle = nullptr;
	30	+ TritonLaunchKernelFn launch_fn = nullptr;
	31	+ std::vector<int> tensor_kinds; // 0=input, 1=output, 2=input_output
	32	+ std::vector<bool> is_constexpr;
	33	+
	34	+ std::string toString() const {
	35	+ std::ostringstream oss;
	36	+ oss << "CompiledKernelEntry{kernel_name=" << kernel_name
	37	+ << ", kernel_func_ptr=" << kernel_func_ptr
	38	+ << ", launcher_so=" << launcher_so_path
	39	+ << ", launch_fn=" << reinterpret_cast<const void*>(launch_fn)
	40	+ << ", num_tensors=" << tensor_kinds.size()
	41	+ << ", num_constexpr=";
	42	+ int ce_count = 0;
	43	+ for (bool b : is_constexpr) if (b) ce_count++;
	44	+ oss << ce_count << "/" << is_constexpr.size();
	45	+ oss << "}";
	46	+ return oss.str();
	47	+ }
	48	+};
	49	+
	50	+class KernelCache {
	51	+public:
	52	+ static KernelCache& instance();
	53	+
	54	+ static uint64_t compute_cache_key(
	55	+ const KernelDescriptor& kernel_desc,
	56	+ const std::vector<ArgInfo>& args);
	57	+
	58	+ std::shared_ptr<CompiledKernelEntry> query(
	59	+ const std::string& kernel_name,
	60	+ uint64_t cache_key) const;
	61	+
	62	+ void store(const std::string& kernel_name,
	63	+ uint64_t cache_key,
	64	+ std::shared_ptr<CompiledKernelEntry> entry);
	65	+
	66	+ void invalidate(const std::string& kernel_name);
	67	+ void clear();
	68	+ void shutdown();
	69	+
	70	+private:
	71	+ KernelCache() = default;
	72	+ std::unordered_map<std::string,
	73	+ std::unordered_map<uint64_t, std::shared_ptr<CompiledKernelEntry>>> cache_;
	74	+ mutable std::shared_mutex mutex_;
	75	+};
	76	+
	77	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/kernel_compiler.h+33-0

@@ -0,0 +1,33 @@
	1	+#pragma once
	2	+#include "kernel_cache.h"
	3	+#include "kernel_registry.h"
	4	+#include "arg_info.h"
	5	+#include <pybind11/pybind11.h>
	6	+#include <vector>
	7	+#include <string>
	8	+#include <functional>
	9	+#include <utility>
	10	+
	11	+namespace triton_runtime {
	12	+
	13	+class KernelCompiler {
	14	+public:
	15	+ static KernelCompiler& instance();
	16	+
	17	+ std::shared_ptr<CompiledKernelEntry> compile(
	18	+ const KernelDescriptor& kernel_desc,
	19	+ const std::vector<ArgInfo>& args,
	20	+ const Grid& grid);
	21	+
	22	+ void shutdown();
	23	+
	24	+private:
	25	+ KernelCompiler() = default;
	26	+
	27	+ pybind11::object do_python_compile(
	28	+ const KernelDescriptor& kernel_desc,
	29	+ const std::vector<ArgInfo>& args,
	30	+ const Grid& grid);
	31	+};
	32	+
	33	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/kernel_registry.h+66-0

@@ -0,0 +1,66 @@
	1	+#pragma once
	2	+#include <string>
	3	+#include <vector>
	4	+#include <memory>
	5	+#include <unordered_map>
	6	+#include <shared_mutex>
	7	+#include <sstream>
	8	+#include <pybind11/pybind11.h>
	9	+#include "status.h"
	10	+
	11	+namespace triton_runtime {
	12	+
	13	+struct KernelDescriptor {
	14	+ std::string kernel_name;
	15	+ std::string python_file;
	16	+ std::string function_name;
	17	+ std::string signature_json;
	18	+ std::string signature_str; // formatted: "def fn(arg: [type], ...)"
	19	+ pybind11::object jit_function;
	20	+ std::vector<bool> is_constexpr;
	21	+ std::vector<std::string> param_names; // params[i].name, for C++ bound_args
	22	+ bool is_autotuned = false;
	23	+
	24	+ std::string toString() const {
	25	+ std::ostringstream oss;
	26	+ oss << "KernelDescriptor{kernel_name=" << kernel_name
	27	+ << ", python_file=" << python_file
	28	+ << ", function_name=" << function_name
	29	+ << ", signature_str=" << signature_str
	30	+ << ", num_params=" << param_names.size()
	31	+ << ", is_autotuned=" << (is_autotuned ? "true" : "false")
	32	+ << "}";
	33	+ return oss.str();
	34	+ }
	35	+};
	36	+
	37	+class KernelRegistry {
	38	+public:
	39	+ static KernelRegistry& instance();
	40	+
	41	+ Status register_kernel(const std::string& python_file,
	42	+ const std::string& function_name,
	43	+ const std::string& kernel_name);
	44	+
	45	+ Status register_kernel_from_dir(const std::string& dir_path);
	46	+
	47	+ std::shared_ptr<KernelDescriptor> lookup(const std::string& kernel_name) const;
	48	+ bool has_kernel(const std::string& kernel_name) const;
	49	+ std::vector<std::string> list_kernels() const;
	50	+ void print_kernel_signature(const std::string& kernel_name) const;
	51	+ Status unregister(const std::string& kernel_name);
	52	+ void shutdown();
	53	+
	54	+private:
	55	+ KernelRegistry() = default;
	56	+ Status register_kernel_impl(const std::string& python_file,
	57	+ const std::string& function_name,
	58	+ const std::string& kernel_name,
	59	+ const pybind11::object& inner_jit_fn,
	60	+ bool is_autotuned);
	61	+
	62	+ std::unordered_map<std::string, std::shared_ptr<KernelDescriptor>> registry_;
	63	+ mutable std::shared_mutex mutex_;
	64	+};
	65	+
	66	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/logging.h+24-0

@@ -0,0 +1,24 @@
	1	+#pragma once
	2	+
	3	+#include <cstdlib>
	4	+#include <cstdio>
	5	+#include <cstring>
	6	+
	7	+namespace triton_runtime {
	8	+
	9	+inline bool is_log_enabled() {
	10	+ static bool enabled = []() {
	11	+ const char* env = std::getenv("LOG_TORCH_TRITON_RUNTIME");
	12	+ return env != nullptr && std::strcmp(env, "1") == 0;
	13	+ }();
	14	+ return enabled;
	15	+}
	16	+
	17	+} // namespace triton_runtime
	18	+
	19	+#define TRT_DEBUG(fmt, ...) \
	20	+ do { \
	21	+ if (::triton_runtime::is_log_enabled()) { \
	22	+ std::fprintf(stderr, "[TRT] " fmt "\n", ##__VA_ARGS__); \
	23	+ } \
	24	+ } while (0)

Aexamples/libtorch_triton_ascend/triton_runtime/include/python_env.h+23-0

@@ -0,0 +1,23 @@
	1	+#pragma once
	2	+#include <string>
	3	+#include <pybind11/pybind11.h>
	4	+
	5	+namespace triton_runtime {
	6	+
	7	+class PythonEnv {
	8	+public:
	9	+ static PythonEnv& instance();
	10	+
	11	+ void ensure_initialized();
	12	+ bool is_initialized() const;
	13	+
	14	+ pybind11::object exec(const std::string& code);
	15	+ pybind11::module import(const std::string& module_name);
	16	+
	17	+private:
	18	+ PythonEnv() = default;
	19	+ bool initialized_ = false;
	20	+ void setup_triton_env();
	21	+};
	22	+
	23	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/status.h+21-0

@@ -0,0 +1,21 @@
	1	+#pragma once
	2	+#include <string>
	3	+
	4	+namespace triton_runtime {
	5	+
	6	+class Status {
	7	+public:
	8	+ static Status OK() { return Status(true, ""); }
	9	+ static Status Error(const std::string& msg) { return Status(false, msg); }
	10	+
	11	+ bool ok() const { return ok_; }
	12	+ const std::string& error_message() const { return msg_; }
	13	+ explicit operator bool() const { return ok_; }
	14	+
	15	+private:
	16	+ Status(bool ok, std::string msg) : ok_(ok), msg_(std::move(msg)) {}
	17	+ bool ok_;
	18	+ std::string msg_;
	19	+};
	20	+
	21	+} // namespace triton_runtime

Aexamples/libtorch_triton_ascend/triton_runtime/include/triton_runtime.h+70-0

@@ -0,0 +1,70 @@
	1	+#pragma once
	2	+
	3	+#include "logging.h"
	4	+#include "grid.h"
	5	+#include "arg_info.h"
	6	+#include "status.h"
	7	+
	8	+#include <string>
	9	+#include <vector>
	10	+#include <torch/torch.h>
	11	+
	12	+namespace triton_runtime {
	13	+
	14	+struct KernelDescriptor;
	15	+
	16	+class TritonRuntime {
	17	+public:
	18	+ static TritonRuntime& instance();
	19	+
	20	+ // ──── Registration Interface ────
	21	+
	22	+ Status register_kernel(const std::string& python_file,
	23	+ const std::string& function_name,
	24	+ const std::string& kernel_name = "");
	25	+
	26	+ Status register_kernel_dir(const std::string& dir_path);
	27	+
	28	+ Status register_kernel_from_dir(const std::string& dir_path);
	29	+
	30	+ // ──── Execution Interface ────
	31	+
	32	+ template<typename... Args>
	33	+ Status run(const std::string& kernel_name,
	34	+ const Grid& grid_spec,
	35	+ Args&&... args) {
	36	+ return run_impl(kernel_name, grid_spec,
	37	+ {ArgInfo(std::forward<Args>(args))...});
	38	+ }
	39	+
	40	+ // ──── Query Interface ────
	41	+
	42	+ bool has_kernel(const std::string& kernel_name) const;
	43	+ std::vector<std::string> list_kernels() const;
	44	+ void print_kernel_signature(const std::string& kernel_name) const;
	45	+
	46	+ // ──── Lifecycle ────
	47	+
	48	+ // Actively clean up all resources (safe to call while Python/CANN is still alive).
	49	+ // If not called, resources will be automatically reclaimed by OS at process exit;
	50	+ // no cleanup during static destruction.
	51	+ void shutdown();
	52	+
	53	+private:
	54	+ TritonRuntime();
	55	+ // ~TritonRuntime() {
	56	+ // shutdown();
	57	+ // }
	58	+
	59	+ // Build a C++ BoundArgs map {param_name: &arg_info, ...} using the
	60	+ // kernel's param_names metadata. Pure C++, no Python.
	61	+ BoundArgs build_bound_args(
	62	+ const KernelDescriptor& desc,
	63	+ const std::vector<ArgInfo>& arg_infos);
	64	+
	65	+ Status run_impl(const std::string& kernel_name,
	66	+ const Grid& grid_spec,
	67	+ std::vector<ArgInfo>&& arg_infos);
	68	+};
	69	+
	70	+} // namespace triton_runtime