Coverage for paintmaskanimator/auto_gpu_regions.py: 8%
237 statements
« prev ^ index » next coverage.py v7.17.0, created at 2026-11-02 04:50 +0000
1"""自動塗りの連結成分と、ラン圧縮した見本の採色をGPUで処理する。"""
2from collections import OrderedDict, defaultdict
3from contextlib import contextmanager
4from hashlib import blake2b
5from threading import RLock
6from time import perf_counter
7from typing import Any
9import numpy as np
11from .logging_setup import get_logger
13log = get_logger(__name__)
16_SOURCE = r'''
17extern "C" __global__ void count_runs(
18 const long long* runs, const int* jobs, const int* params, int nruns,
19 const long long* rows, const int* starts, const int* ends,
20 const int* keys, const int* labels, const unsigned char* flags,
21 int width, int height, int nk, int* counts, long long* out) {
22 int r = blockDim.x * blockIdx.x + threadIdx.x;
23 if (r >= nruns) return;
24 int j = jobs[r]; const int* p = params + j*4;
25 long long y = runs[r*3] + p[0];
26 long long x = runs[r*3+1] + p[1], end = runs[r*3+2] + p[1];
27 if (y < 0 || y >= height) return;
28 x = x < 0 ? 0 : x; end = end > width ? width : end;
29 if (x >= end) return;
30 atomicAdd((unsigned long long*)(out+j*9), (unsigned long long)(end-x));
31 long long lo = rows[y], hi = rows[y+1];
32 while (lo < hi) { long long m=(lo+hi)/2; if(ends[m]<=x) lo=m+1; else hi=m; }
33 for(long long s=lo; s<rows[y+1] && starts[s]<end; ++s) {
34 if ((flags[s]&p[2]) || (flags[s]&p[3])!=p[3]) continue;
35 int n = (ends[s]<end ? ends[s]:end) - (starts[s]>x ? starts[s]:x);
36 atomicAdd(counts+j*nk+keys[s], n);
37 atomicAdd((unsigned long long*)(out+j*9+1), (unsigned long long)n);
38 if(flags[s]&4) atomicExch((unsigned long long*)(out+j*9+6), 1ULL);
39 }
40}
41extern "C" __global__ void winners(const int* counts, const long long* palette,
42 int nk, int nj, int* best, long long* out) {
43 int j=blockDim.x*blockIdx.x+threadIdx.x; if(j>=nj) return;
44 int key=-1, n=0;
45 for(int k=0; k<nk; ++k) if(counts[j*nk+k]>n) { n=counts[j*nk+k]; key=k; }
46 best[j]=key; out[j*9+2]=key<0 ? -1 : palette[key]; out[j*9+3]=n;
47}
48extern "C" __global__ void matched_runs(
49 const long long* runs, const int* jobs, const int* params, int nruns,
50 const long long* rows, const int* starts, const int* ends,
51 const int* keys, const int* labels, const unsigned char* flags,
52 int width, int height, int nl, const int* best, int* touched,
53 long long* lengths, long long* out) {
54 int r=blockDim.x*blockIdx.x+threadIdx.x; if(r>=nruns) return;
55 int j=jobs[r]; const int* p=params+j*4;
56 long long y=runs[r*3]+p[0], x=runs[r*3+1]+p[1], end=runs[r*3+2]+p[1];
57 if(y<0 || y>=height || best[j]<0) return;
58 x=x<0 ? 0:x; end=end>width ? width:end; if(x>=end) return;
59 long long lo=rows[y], hi=rows[y+1];
60 while(lo<hi) { long long m=(lo+hi)/2; if(ends[m]<=x) lo=m+1; else hi=m; }
61 long long n=0;
62 for(long long s=lo; s<rows[y+1] && starts[s]<end; ++s) {
63 if(keys[s]!=best[j] || (flags[s]&p[2]) || (flags[s]&p[3])!=p[3]) continue;
64 n+=(ends[s]<end ? ends[s]:end)-(starts[s]>x ? starts[s]:x);
65 atomicExch(touched+j*nl+labels[s], 1);
66 if(flags[s]&4) atomicExch((unsigned long long*)(out+j*9+5), 1ULL);
67 }
68 lengths[r]=n;
69}
70extern "C" __global__ void finish_samples(
71 const long long* runs, const int* offsets, const int* params, const long long* prefix,
72 const long long* rows, const int* starts, const int* ends, const int* keys,
73 const unsigned char* flags, const long long* sizes, const int* touched,
74 int nl, int nj, const int* best, long long* out) {
75 int j=blockDim.x*blockIdx.x+threadIdx.x; if(j>=nj || best[j]<0) return;
76 long long area=0; for(int k=0; k<nl; ++k) if(touched[j*nl+k]) area+=sizes[k];
77 out[j*9+4]=area;
78 int begin=offsets[j], end=offsets[j+1];
79 long long base=begin ? prefix[begin-1]:0, at=base+out[j*9+3]/2;
80 int lo=begin, hi=end;
81 while(lo<hi) { int m=(lo+hi)/2; if(prefix[m]<=at) lo=m+1; else hi=m; }
82 if(lo>=end) return;
83 int r=lo; const int* p=params+j*4;
84 long long skip=at-(r ? prefix[r-1]:0);
85 long long y=runs[r*3]+p[0], x=runs[r*3+1]+p[1], stop=runs[r*3+2]+p[1];
86 x=x<0 ? 0:x;
87 lo=rows[y]; hi=rows[y+1];
88 while(lo<hi) { int m=(lo+hi)/2; if(ends[m]<=x) lo=m+1; else hi=m; }
89 for(long long s=lo; s<rows[y+1] && starts[s]<stop; ++s) {
90 if(keys[s]!=best[j] || (flags[s]&p[2]) || (flags[s]&p[3])!=p[3]) continue;
91 long long left=starts[s]>x ? starts[s]:x;
92 long long n=(ends[s]<stop ? ends[s]:stop)-left;
93 if(skip<n) { out[j*9+7]=y; out[j*9+8]=left+skip; return; }
94 skip-=n;
95 }
96}
97'''
100def mask_runs(filled):
101 runs = getattr(filled, "_native_runs", None)
102 if runs is None:
103 padded = np.pad(filled.local, ((0, 0), (1, 1)))
104 ys, xs = np.nonzero(padded[:, 1:] != padded[:, :-1])
105 runs = np.ascontiguousarray(np.column_stack((
106 ys[::2] + filled.bounds[1], xs[::2] + filled.bounds[0], xs[1::2] + filled.bounds[0],
107 )), dtype=np.int64).tobytes()
108 filled._native_runs = runs
109 return np.frombuffer(runs, dtype=np.int64).reshape(-1, 3)
112class RegionGPU:
113 """計画内だけのGPUキャッシュ。採色結果のみCPUへ戻し、不変の見本を再利用する。"""
115 def __init__(self, cp, *, max_bytes=128 * 1024 * 1024, capacity=16, mode="on"):
116 self.cp = cp
117 self.device = cp.cuda.Device().id
118 self.pool = cp.cuda.MemoryPool()
119 self.lock = RLock()
120 self.guides = OrderedDict()
121 self.results = OrderedDict()
122 self.max_bytes = max_bytes
123 self.capacity = capacity
124 self.failed = False
125 self.error = None
126 self.sample_count = self.sample_batches = self.index_count = 0
127 self.sample_seconds = self.index_seconds = 0.0
128 self.guide_uploads = 0
129 self.result_capacity = 4096
130 self._pinned = 0
131 from .auto_routing import PipelinePolicy
132 self.mode = mode
133 self.policy = PipelinePolicy()
134 self.cpu_samples = self.cpu_indices = 0
135 self.device_shift_samples = 0
136 self.kernels = {name: cp.RawKernel(_SOURCE, name) for name in (
137 "count_runs", "winners", "matched_runs", "finish_samples",
138 )}
140 def close(self):
141 with self.lock:
142 self.guides.clear()
143 self.results.clear()
144 try:
145 with self.cp.cuda.Device(self.device):
146 self.pool.free_all_blocks()
147 except Exception as error:
148 self._failure(error)
150 def _failure(self, error):
151 self.error = str(error)
152 self.failed = True
153 log.warning("GPU region processing failed; continuing on CPU: %s", self.error)
155 def index(self, passable):
156 """4近傍の連結成分と外接矩形。背景番号0は種として使わない。"""
157 if self.failed or passable.size == 0:
158 return None
159 policy_key = ('bucket', passable.size.bit_length(),
160 int(np.count_nonzero(passable[:, 1:] != passable[:, :-1])).bit_length())
161 route = self.policy.choose(policy_key) if self.mode == 'auto' else True
162 cpu_result, cpu_seconds = None, 0.0
163 if route is not True:
164 from .bucket_fill import _RegionIndex
165 started = perf_counter()
166 cpu = _RegionIndex(passable)
167 cpu_result = cpu.labels, cpu.x0, cpu.y0, cpu.x1, cpu.y1
168 cpu_seconds = perf_counter() - started
169 if route is False:
170 self.cpu_indices += 1
171 return cpu_result
172 try:
173 with self.lock, self.cp.cuda.Device(self.device), self.cp.cuda.using_allocator(self.pool.malloc):
174 from importlib import import_module
175 label = import_module("cupyx.scipy.ndimage").label
176 started = perf_counter()
177 cp = self.cp
178 labeled: Any = label(cp.asarray(passable))
179 labels, count = labeled
180 height, width = passable.shape
181 change = cp.ones(labels.shape, dtype=bool)
182 change[:, 1:] = labels[:, 1:] != labels[:, :-1]
183 ys, xs = cp.nonzero(change)
184 ends = cp.r_[xs[1:], width]
185 ends[cp.r_[ys[1:] != ys[:-1], True]] = width
186 kinds = labels[ys, xs]
187 x0, y0 = cp.full(count + 1, width, dtype=cp.int64), cp.full(count + 1, height, dtype=cp.int64)
188 x1, y1 = cp.zeros(count + 1, dtype=cp.int64), cp.zeros(count + 1, dtype=cp.int64)
189 cp.minimum.at(x0, kinds, xs)
190 cp.minimum.at(y0, kinds, ys)
191 cp.maximum.at(x1, kinds, ends)
192 cp.maximum.at(y1, kinds, ys + 1)
193 result = tuple(cp.asnumpy(array) for array in (labels, x0, y0, x1, y1))
194 self.index_count += 1
195 elapsed = perf_counter() - started
196 self.index_seconds += elapsed
197 if route is None:
198 use_gpu = self.policy.record(policy_key, cpu_seconds, elapsed)
199 if not use_gpu:
200 self.cpu_indices += 1
201 return cpu_result
202 return result
203 except Exception as error:
204 self._failure(error)
205 return None
207 def _guide(self, guide):
208 if guide in self.guides:
209 result = self.guides.pop(guide)
210 self.guides[guide] = result
211 return result
212 cp = self.cp
213 initial = guide.initial if guide.initial is not None else guide
214 flags = (guide.boundary.astype(np.uint8) | initial.boundary.astype(np.uint8) * 2
215 | (guide.open | guide.undecided).astype(np.uint8) * 4)
216 confirmed = getattr(guide, "confirmed", None)
217 flags |= 8 if confirmed is None else confirmed.astype(np.uint8) * 8
218 change: np.ndarray = np.ones(flags.shape, dtype=bool)
219 change[:, 1:] = ((guide.values[:, 1:] != guide.values[:, :-1])
220 | (guide.regions.labels[:, 1:] != guide.regions.labels[:, :-1])
221 | (flags[:, 1:] != flags[:, :-1]))
222 ys, xs = np.nonzero(change)
223 ends = np.r_[xs[1:], flags.shape[1]]
224 ends[np.r_[ys[1:] != ys[:-1], True]] = flags.shape[1]
225 values = guide.values[ys, xs]
226 palette, keys = np.unique(values, return_inverse=True)
227 arrays = (np.searchsorted(ys, np.arange(flags.shape[0] + 1)).astype(np.int64),
228 xs.astype(np.int32), ends.astype(np.int32), keys.astype(np.int32),
229 guide.regions.labels[ys, xs].astype(np.int32), flags[ys, xs],
230 guide.regions.sizes.astype(np.int64), palette.astype(np.int64))
231 result = tuple(cp.asarray(array) for array in arrays)
232 self.guide_uploads += 1
233 size = sum(array.nbytes for array in result)
234 while self.guides and (len(self.guides) >= self.capacity or
235 sum(sum(a.nbytes for a in item) for item in self.guides.values()) + size > self.max_bytes):
236 self.guides.popitem(last=False)
237 if size <= self.max_bytes:
238 self.guides[guide] = result
239 return result
241 @staticmethod
242 def key(guide, filled, shift, ignore, require, by_label):
243 mask_key = getattr(filled, "_gpu_sample_key", None)
244 if mask_key is None:
245 mask_key = blake2b(mask_runs(filled).tobytes(), digest_size=16).digest()
246 filled._gpu_sample_key = mask_key
247 return guide, mask_key, tuple(shift), ignore, require, by_label
249 def cached(self, key):
250 return self.results.get(key)
252 @contextmanager
253 def phase(self):
254 # 再試行中の依存結果をLRUで失うと、要求が入れ替わり続ける場合がある。
255 # 1段階の採色がすべて完了するまで保持し、完了・例外時に上限へ戻す。
256 self._pinned += 1
257 try:
258 yield
259 finally:
260 self._pinned -= 1
261 if not self._pinned:
262 while len(self.results) > self.result_capacity:
263 self.results.popitem(last=False)
265 def sample_many(self, requests):
266 """同じ見本の採色をまとめ、画素ではなくランの重なりをGPUで集計する。"""
267 if self.failed:
268 return
269 try:
270 with self.lock, self.cp.cuda.Device(self.device), self.cp.cuda.using_allocator(self.pool.malloc):
271 groups = defaultdict(list)
272 for key, request in requests.items():
273 groups[(request[0], request[-1])].append((key, request))
274 for (guide, by_label), items in groups.items():
275 nl = len(guide.regions.sizes)
276 capacity = min(64, max(1, 16 * 1024 * 1024 // (8 * nl)))
277 batches, batch, run_count = [], [], 0
278 for item in items:
279 count = len(mask_runs(item[1][1]))
280 if batch and (len(batch) >= capacity or run_count + count > 262144):
281 batches.append(batch)
282 batch, run_count = [], 0
283 batch.append(item)
284 run_count += count
285 if batch:
286 batches.append(batch)
287 for batch in batches:
288 work = sum(len(mask_runs(item[1][1])) for item in batch)
289 policy_key = ('sample', len(batch).bit_length(), work.bit_length(), by_label,
290 guide in self.guides)
291 route = self.policy.choose(policy_key) if self.mode == 'auto' else True
292 cpu_answers, cpu_seconds = None, 0.0
293 if route is not True:
294 from . import native
295 started = perf_counter()
296 cpu_answers = [native.sample_region(r[0], r[1], r[2], ignore=r[3],
297 require=r[4], by_label=r[5])
298 for _key, r in batch]
299 cpu_seconds = perf_counter() - started
300 if any(answer is None for answer in cpu_answers):
301 cpu_answers, route = None, True
302 if route is False:
303 answers = cpu_answers
304 self.cpu_samples += len(batch)
305 else:
306 started = perf_counter()
307 device = self._guide(guide)
308 answers = self._sample_batch(guide, device, batch, by_label)
309 elapsed = perf_counter() - started
310 self.sample_seconds += elapsed
311 self.sample_count += len(batch)
312 self.sample_batches += 1
313 if route is None:
314 valid = answers == cpu_answers
315 self.policy.record(policy_key, cpu_seconds, elapsed, valid=valid)
316 if not valid:
317 self._failure(RuntimeError('GPU sampling differs from CPU'))
318 answers = cpu_answers
319 self.cpu_samples += len(batch)
320 assert answers is not None
321 for (key, _request), answer in zip(batch, answers):
322 # 大きなマスクをキャッシュに保持せず、圧縮ランの指紋で照合する。
323 self.results[key] = answer
324 while not self._pinned and len(self.results) > max(self.result_capacity, len(requests)):
325 self.results.popitem(last=False)
326 except Exception as error:
327 self._failure(error)
329 def _sample_batch(self, guide, device, batch, by_label):
330 cp = self.cp
331 rows, starts, ends, color_keys, labels, flags, sizes, palette = device
332 keys = labels if by_label else color_keys
333 palette = cp.arange(len(sizes), dtype=cp.int64) if by_label else palette
334 masks = [mask_runs(request[1]) for _key, request in batch]
335 offsets = np.r_[0, np.cumsum([len(mask) for mask in masks])].astype(np.int32)
336 runs = cp.asarray(np.concatenate(masks))
337 jobs = cp.asarray(np.repeat(np.arange(len(batch), dtype=np.int32), np.diff(offsets)))
338 params = cp.asarray(np.array([
339 (0, 0, r[3], r[4]) if hasattr(r[2], 'device') else (r[2][0], r[2][1], r[3], r[4])
340 for _key, r in batch], dtype=np.int32))
341 device_jobs = [(index, r[2].device) for index, (_key, r) in enumerate(batch)
342 if hasattr(r[2], 'device')]
343 if device_jobs:
344 indices = cp.asarray(np.array([index for index, _shift in device_jobs], dtype=np.int32))
345 params[indices, :2] = cp.stack([shift for _index, shift in device_jobs]).astype(cp.int32)
346 self.device_shift_samples += len(device_jobs)
347 nj, nr, nk, nl = len(batch), len(runs), len(palette), len(sizes)
348 counts = cp.zeros((nj, nk), dtype=cp.int32)
349 output = cp.zeros((nj, 9), dtype=cp.int64)
350 best = cp.empty(nj, dtype=cp.int32)
351 touched = cp.zeros((nj, nl), dtype=cp.int32)
352 lengths = cp.zeros(nr, dtype=cp.int64)
353 height, width = guide.values.shape
354 scalar = lambda value: np.int32(value)
355 common = (runs, jobs, params, scalar(nr), rows, starts, ends, keys, labels, flags,
356 scalar(width), scalar(height))
357 if nr:
358 self.kernels["count_runs"](((nr + 255) // 256,), (256,),
359 (*common, scalar(nk), counts, output))
360 self.kernels["winners"](((nj + 127) // 128,), (128,),
361 (counts, palette, scalar(nk), scalar(nj), best, output))
362 if nr:
363 self.kernels["matched_runs"](((nr + 255) // 256,), (256,),
364 (*common, scalar(nl), best, touched, lengths, output))
365 prefix = cp.cumsum(lengths)
366 self.kernels["finish_samples"](((nj + 127) // 128,), (128,), (
367 runs, cp.asarray(offsets), params, prefix, rows, starts, ends, keys, flags, sizes,
368 touched, scalar(nl), scalar(nj), best, output,
369 ))
370 answers = cp.asnumpy(output)
371 return [(int(a), int(b), None if c < 0 else int(c), int(d), int(e), bool(f), bool(g), int(h), int(i))
372 for a, b, c, d, e, f, g, h, i in answers]