Line data Source code
1 : /**
2 : * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3 : * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 : * CANN Open Software License Agreement Version 2.0 (the "License").
5 : * Please refer to the License for details. You may not use this file except in compliance with the License.
6 : * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7 : * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8 : * See LICENSE in the root of the software repository for the full text of the License.
9 : */
10 :
11 : #ifndef AICPU_SHARDER_H
12 : #define AICPU_SHARDER_H
13 :
14 : #include <mutex>
15 : #include <atomic>
16 : #include <functional>
17 : #include <vector>
18 : #include <queue>
19 :
20 : namespace aicpu {
21 : using Closure = std::function<void()>;
22 : using SharderWork = std::function<void(int64_t, int64_t)>;
23 : using RandomKernelScheduler = std::function<uint32_t(const aicpu::Closure &task)>;
24 : using SplitKernelScheduler = std::function<uint32_t(const uint32_t parallelId, const int64_t shardNum,
25 : const std::queue<Closure> &queue)>;
26 : using SplitKernelGetProcesser = std::function<bool()>;
27 :
28 : class SharderNonBlock {
29 : public:
30 : /**
31 : * Get the unique object of this class
32 : */
33 : static SharderNonBlock &GetInstance();
34 :
35 : /**
36 : * Register schedule callback function, doTask function and cpu core number
37 : * called by compute process
38 : * @param cpuCoreNum aicpu core number
39 : * @param randomKernelScheduler random kernel event submit and task enqueue
40 : * @param splitKernelScheduler random kernel event submit and task add
41 : * @param splitKernelScheduler get and process split kernel for main thread
42 : */
43 : void Register(const uint32_t cpuCoreNum, const RandomKernelScheduler &randomKernelScheduler,
44 : const SplitKernelScheduler &splitKernelScheduler,
45 : const SplitKernelGetProcesser &splitKernelGetProcesser);
46 :
47 : /**
48 : * Shards the "total" unit of work refer "perUintSize"
49 : * @param total Total unit of work
50 : * @param perUnitSize Minimum shard unit
51 : * @param work should be a callable taking (int64, int64) arguments.
52 : work(start, limit) computes the work units from [start, limit),
53 : i.e., [start, limit) is a shard.
54 : */
55 : void ParallelFor(const int64_t total, const int64_t perUnitSize, const SharderWork &work);
56 :
57 : /**
58 : * Shards the unit of work refer for hash
59 : * @param total, Total unit of work
60 : * @param cpuNums Number of cpu cores
61 : * @param work should be a callable taking (int64, int64) arguments.
62 : work(cur, cpuNums) computes the work units with input hash with (cpuNums-1) equals cur,
63 : i.e. specially used by parallel unique op
64 : */
65 : void ParallelForHash(const int64_t total, const int64_t cpuNums, const SharderWork &work);
66 :
67 : /**
68 : * Schedule a task use schedule function registered by compute process,
69 : * note that the task will actually executed asynchronously
70 : * @param closure Closure function with nothrow
71 : */
72 : void Schedule(const Closure &aicpuClosure);
73 :
74 : /**
75 : * Get CPU number
76 : * @param None
77 : * @return CPU number
78 : */
79 : uint32_t GetCPUNum();
80 :
81 : private:
82 : SharderNonBlock();
83 5 : ~SharderNonBlock() = default;
84 :
85 : SharderNonBlock(const SharderNonBlock &) = delete;
86 : SharderNonBlock &operator = (const SharderNonBlock &) = delete;
87 : SharderNonBlock(SharderNonBlock &&) = delete;
88 : SharderNonBlock &operator = (SharderNonBlock &&) = delete;
89 :
90 : void DoTaskItself(const uint32_t parallelId, std::atomic<int64_t> &cpuNumCounter,
91 : const int64_t shardNum) const;
92 :
93 : /**
94 : * Calculate how many times, which ceiled, "x" is "base".
95 : * i.e., x is 1, base is 2, this function will return 1
96 : * @param x An integral
97 : * @param base An integral as base when cal multiple
98 : * @return ceiled multiple
99 : */
100 : inline int64_t CeilMultiple(const int64_t x, const int64_t base) const;
101 :
102 : /**
103 : * Shards the "total" unit of work refer "perUintSize"
104 : * @param total Total unit of work
105 : * @param shardNum parralle number
106 : * @param blockSize Minimum shard unit
107 : * @param work should be a callable taking (int64, int64) arguments.
108 : work(start, limit) computes the work units from [start, limit),
109 : i.e., [start, limit) is a shard.
110 : */
111 : void ExecuteParallelFor(const int64_t total, const int64_t shardNum,
112 : const int64_t blockSize, const SharderWork &work, const uint32_t parallelId);
113 :
114 : void ExecuteParallelForHash(const int64_t total, const int64_t cpuNums, const SharderWork &work,
115 : const uint32_t parallelId);
116 :
117 : private:
118 : uint32_t cpuCoreNum_; // aicpu core number
119 : RandomKernelScheduler randomKernelScheduler_;
120 : SplitKernelScheduler splitKernelScheduler_;
121 : SplitKernelGetProcesser splitKernelGetProcesser_;
122 : std::atomic<uint32_t> parallelId_; // the id for parallel run kernel
123 : std::mutex parallelIdMutex_;
124 : }; // SharderNonBlock
125 : } // namespace aicpu
126 :
127 : extern "C" {
128 : /**
129 : * Shards the "total" unit of work refer "perUintSize"
130 : * @param total Total unit of work
131 : * @param perUnitSize Minimum shard unit
132 : * @param work should be a callable taking (int64, int64) arguments.
133 : work(start, limit) computes the work units from [start, limit),
134 : i.e., [start, limit) is a shard.
135 : */
136 : __attribute__((visibility("default"))) void ParallelFor(int64_t total, int64_t perUnitSize,
137 : const aicpu::SharderWork &work);
138 :
139 : /**
140 : * Get CPU number
141 : * @param None
142 : * @return CPU number
143 : */
144 : __attribute__((visibility("default"))) uint32_t GetCPUNum();
145 : }
146 :
147 : #endif // AICPU_SHARDER_H_
|
