From 23e3721481292aa9384f027520517ebf7851cd8d Mon Sep 17 00:00:00 2001 From: b00874339 Date: Tue, 5 Aug 2025 21:52:23 +0800 Subject: [PATCH] Feature: Add kupl backend & examples --- csrc/utils/parallel.h | 84 +++++++++++++++++++ kpex_example/__init__.py | 1 + kpex_example/csrc/kpex_example.cpp | 11 +++ kpex_example/csrc/kupl_example/bind.h | 15 ++++ kpex_example/csrc/kupl_example/kupl_example.h | 11 +++ .../csrc/kupl_example/kupl_example_pf.cpp | 44 ++++++++++ kpex_example/frontend.py | 0 kpex_example/kupl_example/__init__.py | 1 + kpex_example/kupl_example/kupl_example.py | 16 ++++ setup.py | 27 ++++++ 10 files changed, 210 insertions(+) create mode 100644 csrc/utils/parallel.h create mode 100644 kpex_example/__init__.py create mode 100644 kpex_example/csrc/kpex_example.cpp create mode 100644 kpex_example/csrc/kupl_example/bind.h create mode 100644 kpex_example/csrc/kupl_example/kupl_example.h create mode 100644 kpex_example/csrc/kupl_example/kupl_example_pf.cpp create mode 100644 kpex_example/frontend.py create mode 100644 kpex_example/kupl_example/__init__.py create mode 100644 kpex_example/kupl_example/kupl_example.py diff --git a/csrc/utils/parallel.h b/csrc/utils/parallel.h new file mode 100644 index 0000000..660863b --- /dev/null +++ b/csrc/utils/parallel.h @@ -0,0 +1,84 @@ +// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved + +#pragma once + +#include +#include +#include +#include "kupl.h" + + +namespace kpex { +namespace { +inline int64_t divup(int64_t x, int64_t y) { + return (x + y - 1) / y; +} + +template +struct parallel_for_function_args { + int64_t begin; + int64_t end; + int64_t grain_size; + const F& f; +}; + +template +static void parallel_for_function(kupl_nd_range_t *nd_range, void *args, int tid, int tnum) +{ + auto data = (parallel_for_function_args *)args; + int64_t begin = data->begin; + int64_t end = data->end; + int64_t grain_size = data->grain_size; + const F& f = data->f; + int64_t num_threads = tnum; + if (grain_size > 0) { + num_threads = std::min(num_threads, divup((end - begin), grain_size)); + } + + int64_t chunk_size = divup((end - begin), num_threads); + int64_t begin_tid = begin + tid * chunk_size; + if (begin_tid < end) { + f(begin_tid, std::min(end, chunk_size + begin_tid)); + } +} + +template +struct parallel_function_args { + const F& f; +}; + +template +static void parallel_function(kupl_nd_range_t *nd_range, void *args, int tid, int tnum) +{ + auto data = (parallel_for_function_args *)args; + const F& f = data->f; + f(tid); +} +} + +template +inline void parallel_for(int64_t begin, int64_t end, int64_t grain_size, const F& f) { + int num_executors = kupl_get_num_executors(); + kupl_parallel_for_desc_t desc = { + .range = NULL, + .concurrency = num_executors, + .egroup = NULL, + .policy = KUPL_LOOP_POLICY_STATIC + }; + parallel_for_function_args args = {begin, end, grain_size, f}; + kupl_parallel_for(&desc, parallel_for_function, &args); +} + +template +inline void parallel(int num_threads, const F& f) { + kupl_parallel_for_desc_t desc = { + .range = NULL, + .concurrency = num_threads, + .egroup = NULL, + .policy = KUPL_LOOP_POLICY_STATIC + }; + parallel_function_args args = {f}; + kupl_parallel_for(&desc, parallel_function, &args); +} + +} // kpex diff --git a/kpex_example/__init__.py b/kpex_example/__init__.py new file mode 100644 index 0000000..f5ea76b --- /dev/null +++ b/kpex_example/__init__.py @@ -0,0 +1 @@ +from . import kupl_example \ No newline at end of file diff --git a/kpex_example/csrc/kpex_example.cpp b/kpex_example/csrc/kpex_example.cpp new file mode 100644 index 0000000..a0b5113 --- /dev/null +++ b/kpex_example/csrc/kpex_example.cpp @@ -0,0 +1,11 @@ +// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved + +#include "kupl_example/bind.h" +#include "utils/memory.h" +#include + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + kupl_example::bind(m); + m.def("device", &kpex::device); +} \ No newline at end of file diff --git a/kpex_example/csrc/kupl_example/bind.h b/kpex_example/csrc/kupl_example/bind.h new file mode 100644 index 0000000..3d8542a --- /dev/null +++ b/kpex_example/csrc/kupl_example/bind.h @@ -0,0 +1,15 @@ +// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved + +#pragma once + +#include + +#include "kupl_example.h" + +namespace kupl_example { +inline void bind(pybind11::module &m) +{ + auto submodule = m.def_submodule("kupl_example"); + submodule.def("test_kupl_parallel_example", &test_kupl_parallel_example); +} +} \ No newline at end of file diff --git a/kpex_example/csrc/kupl_example/kupl_example.h b/kpex_example/csrc/kupl_example/kupl_example.h new file mode 100644 index 0000000..270b0bc --- /dev/null +++ b/kpex_example/csrc/kupl_example/kupl_example.h @@ -0,0 +1,11 @@ +// Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved + +#pragma once + +#include + +namespace kupl_example { +int test_kupl_parallel_example(); +} // namespace kupl_example + + diff --git a/kpex_example/csrc/kupl_example/kupl_example_pf.cpp b/kpex_example/csrc/kupl_example/kupl_example_pf.cpp new file mode 100644 index 0000000..26e1cc3 --- /dev/null +++ b/kpex_example/csrc/kupl_example/kupl_example_pf.cpp @@ -0,0 +1,44 @@ +#include +#include +#include +#include "utils/parallel.h" +#include "kupl_example.h" + +namespace kupl_example { +void test_kupl_parallel() +{ + std::cout << "test_kupl_parallel " << std::endl; + int num_threads = kupl_get_num_executors(); + int tid = kupl_get_executor_num(); + kpex::parallel(num_threads, [&](int tid) { + std::cout << "KUPL Thread " << tid << std::endl; + }); +} + + +void test_kupl_parallel_for() +{ + std::cout << "test_kupl_parallel_for " << std::endl; + int out[1024] = {0}; + kpex::parallel_for(0, 1024, 1, [&](int64_t start, int64_t end) { + for (int64_t i = start; i < end; ++i) { + out[i] = i * i; + } + }); + for (int64_t i = 0; i < 1024; ++i) { + assert(out[i] == i * i); + } +} + +int test_kupl_parallel_example() +{ + test_kupl_parallel(); + + std::cout << "===========================" << std::endl; + sleep(1); + + test_kupl_parallel_for(); + + return 0; +} +} // namespace kupl_example diff --git a/kpex_example/frontend.py b/kpex_example/frontend.py new file mode 100644 index 0000000..e69de29 diff --git a/kpex_example/kupl_example/__init__.py b/kpex_example/kupl_example/__init__.py new file mode 100644 index 0000000..f5ea76b --- /dev/null +++ b/kpex_example/kupl_example/__init__.py @@ -0,0 +1 @@ +from . import kupl_example \ No newline at end of file diff --git a/kpex_example/kupl_example/kupl_example.py b/kpex_example/kupl_example/kupl_example.py new file mode 100644 index 0000000..bf099e5 --- /dev/null +++ b/kpex_example/kupl_example/kupl_example.py @@ -0,0 +1,16 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved + +import copy +import time +import types + +import torch +from torch import nn +import torch.distributed as dist +import kpex._C_example as kernel + +def kupl_parallel_test_example(): + out = kernel.kupl_example.test_kupl_parallel_example() + return out + + diff --git a/setup.py b/setup.py index 5b21ca5..81f511b 100644 --- a/setup.py +++ b/setup.py @@ -17,9 +17,18 @@ sources = [ "csrc/tpp/alphafold/*.cpp", ] +sources_example = [ + "csrc/utils/*.cpp", + "kpex_example/csrc/*.cpp", + "kpex_example/csrc/kupl_example/*.cpp", +] + sources = [j for i in sources for j in glob.glob(i)] +sources_example = [j for i in sources_example for j in glob.glob(i)] extra_compile_args = [] include_dirs = [(root / "csrc").as_posix()] +include_dirs_example = [(root / "csrc").as_posix()] +include_dirs_example += [(root / "kpex_example/csrc").as_posix()] extra_compile_args += ["-fopenmp"] if os.environ["KPEX_BUILD_TYPE"] == "release": @@ -38,6 +47,15 @@ if KUTACC_ROOT: libraries = ["kutacc"] extra_compile_args += [f"-Wl, -rpath={KUTACC_ROOT}/lib"] +library_dirs_example = [] +libraries_example = [] +KUPL_ROOT = os.environ.get("KUPL_ROOT", None) +if KUPL_ROOT: + library_dirs_example += [f"{KUPL_ROOT}/include"] + libraries_example += [f"{KUPL_ROOT}/lib"] + libraries_example = ["kupl"] + extra_compile_args += [f"-Wl, -rpath={KUPL_ROOT}/lib"] + setup( name="kunpeng-pytorch-extension", version="0.0.1", @@ -51,6 +69,15 @@ setup( libraries=libraries, library_dirs=library_dirs, extra_compile_args={"cxx": extra_compile_args}, + ), + + cpp_extension.CppExtension( + name="kpex._C_example", + sources=sources_example, + include_dirs=include_dirs_example, + libraries=libraries_example, + library_dirs=libraries_example, + extra_compile_args={"cxx": extra_compile_args}, ) ], cmdclass={"build_ext": cpp_extension.BuildExtension} -- Gitee