From 743c0e6da596b786a0589d0a87d05e53afa74eda Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 27 Jan 2026 15:22:42 +0900 Subject: [PATCH] avoid buffer size overflow --- cuslines/cuda_python/cu_propagate_seeds.py | 5 +++-- cuslines/cuda_python/cu_tractography.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/cuslines/cuda_python/cu_propagate_seeds.py b/cuslines/cuda_python/cu_propagate_seeds.py index f9a401a..c1d8a8b 100644 --- a/cuslines/cuda_python/cu_propagate_seeds.py +++ b/cuslines/cuda_python/cu_propagate_seeds.py @@ -1,4 +1,5 @@ import numpy as np +import math import gc from cuda.bindings import runtime from cuda.bindings.runtime import cudaMemcpyKind @@ -240,7 +241,7 @@ def get_buffer_size(self): lens = self.sline_lens[ii] for jj in range(self.nSlines[ii]): buffer_size += lens[jj] * 3 * REAL_SIZE - return buffer_size + return math.ceil(buffer_size / MEGABYTE) def as_generator(self): def _yield_slines(): @@ -256,4 +257,4 @@ def _yield_slines(): return _yield_slines() def as_array_sequence(self): - return ArraySequence(self.as_generator(), self.get_buffer_size() // MEGABYTE) + return ArraySequence(self.as_generator(), self.get_buffer_size()) diff --git a/cuslines/cuda_python/cu_tractography.py b/cuslines/cuda_python/cu_tractography.py index 9c24cd7..7064cc9 100644 --- a/cuslines/cuda_python/cu_tractography.py +++ b/cuslines/cuda_python/cu_tractography.py @@ -248,7 +248,7 @@ def generate_sft(self, seeds, ref_img): seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].shape[0] ) array_sequence = ArraySequence( - (item for gen in generators for item in gen), buffer_size // MEGABYTE + (item for gen in generators for item in gen), buffer_size ) return StatefulTractogram(array_sequence, ref_img, Space.VOX)