initial commit

This commit is contained in:
2026-08-07 15:56:42 +09:00
commit 91ad25aca9
1012 changed files with 159314 additions and 0 deletions
+552
View File
@@ -0,0 +1,552 @@
% AX_SUP_PARTIAL distributed (multiGPU) forward projector that allows to split the full volume into smaller pieces
% this allows to solve datasets much larger than memory of used GPU
% or spread calculations over several GPUs
%
% projData_all = Ax_sup_partial(volData, cfg, vectors, split, varargin)
%
% Inputs:
% **volData - array Nx x Ny x Nz of projected volume
% **cfg - config structure generated by ASTRA_initialize
% **vectors - orientation of projections generated by ASTRA_initialize
% **split - 3 or 4 elements vector, [split X, split Y, split Z, split angle ]
% **deformation_fields: 3x1 cell contaning 3D arrays of local deformation of the object
% **verbose - verbose <= 0 : quiet, verbose : standard info , verbose = 2: debug
% **use_shared_memory - true - share data between processed by shared memory, false = use matlab parfor distribution
% **max_memory_blocks - maximal size of used share memory memory
% **varargin - for additional parameters see the code and als the astra.Ax_partial function
% *returns*
% ++projData_all - projection of the volData
%
%
% recompile commands
% (Linux, GCC 4.8.5) mexcuda -outdir private +astra/ASTRA_GPU_wrapper/ASTRA_GPU_wrapper.cu +astra/ASTRA_GPU_wrapper/util3d.cu +astra/ASTRA_GPU_wrapper/par3d_fp.cu +astra/ASTRA_GPU_wrapper/par3d_bp.cu
% (Windows) mexcuda -outdir private ASTRA_GPU_wrapper\ASTRA_GPU_wrapper.cu ASTRA_GPU_wrapper\util3d.cu ASTRA_GPU_wrapper\par3d_fp.cu ASTRA_GPU_wrapper\par3d_bp.cu
%*-----------------------------------------------------------------------*
%|                                                                       |
%|  Except where otherwise noted, this work is licensed under a          |
%|  Creative Commons Attribution-NonCommercial-ShareAlike 4.0            |
%|  International (CC BY-NC-SA 4.0) license.                             |
%|                                                                       |
%|  Copyright (c) 2017 by Paul Scherrer Institute (http://www.psi.ch)    |
%|                                                                       |
%|      Author: CXS group, PSI  |
%*-----------------------------------------------------------------------*
% You may use this code with the following provisions:
%
% If the code is fully or partially redistributed, or rewritten in another
% computing language this notice should be included in the redistribution.
%
% If this code, or subfunctions or parts of it, is used for research in a
% publication or if it is fully or partially rewritten for another
% computing language the authors and institution should be acknowledged
% in written form in the publication: Data processing was carried out
% using the cSAXS matlab package developed by the CXS group,
% Paul Scherrer Institut, Switzerland.
% Variations on the latter text can be incorporated upon discussion with
% the CXS group if needed to more specifically reflect the use of the package
% for the published work.
%
% A publication that focuses on describing features, or parameters, that
% are already existing in the code should be first discussed with the
% authors.
%
% This code and subroutines are part of a continuous development, they
% are provided as they are without guarantees or liability on part
% of PSI or the authors. It is the user responsibility to ensure its
% proper use and the correctness of the results.
function projData_all = Ax_sup_partial(volData, cfg, vectors, split, varargin)
import utils.*
import math.*
par = inputParser;
par.addOptional('deformation_fields', {}) % deformation_fields: 3x1 cell contaning 3D arrays of local deformation of the object
par.addOptional('GPU', []) % list of GPUs to be used in reconstruction
par.addOptional('split_sub', 1) % splitting of the sub block on smaller tasks in the Atx_partial method , 1 == no splitting , otherwise [split_x,split_y,split_z,split_angles]
par.addOptional('verbose', 1) % verbose <= 0 : quiet, verbose : standard info , verbose = 2: debug
par.addOptional('use_shared_memory', []) % true - share data between processed by shared memory, false = use matlab parfor distribution
par.addOptional('max_memory_blocks', utils.check_available_memory*1e6/4) % maximal size of used share memory memory
par.parse(varargin{:})
r = par.Results;
if isscalar(split)
split = split .* ones(1,3);
end
if isempty(r.use_shared_memory)
r.use_shared_memory = length(r.GPU) > 1;
end
%% if not splitting on this level is requirested, continue to tomo.Ax_partial
if all(split(1:3) == 1)
projData_all = astra.Ax_partial(volData, cfg, vectors, r.split_sub, ...
'GPU', r.GPU, 'deformation_fields', r.deformation_fields, 'verbose', r.verbose);
return
end
%% otherwise do proper checking of all inputs before splitting
if ~(isa(volData, 'gpuArray') && strcmp(classUnderlying(volData), 'single')) && ...
~isa(volData, 'single')
error('Only single precision input array supported')
end
volData = gather(volData);
if ismatrix(volData)
% only if volData is 2D array !!
split = split([1, min(2,end)]);
assert(all(size(volData)==[cfg.iVolX,cfg.iVolY]), 'Wrong inputs size')
else
assert(all(size(volData)==[cfg.iVolX,cfg.iVolY,cfg.iVolZ]), 'Wrong inputs size')
end
if isempty(r.GPU)
gpu = gpuDevice;
r.GPU = gpu.Index;
end
N_GPU = length(r.GPU);
% make it equaly splitable among the GPUs
split(3) = ceil(max(split(3), ceil(prod(split(1:3)) / N_GPU) * N_GPU) / prod(split(1:2)));
if length(split) == 3
split(4) = 1;
end
%% backprojector that allows to split the full volume into smaller pieces
cfg.iProjAngles = size(vectors,1);
assert(cfg.iProjAngles > 1, 'Not supported <=1 angles')
assert(all(size(vectors,2)==12), 'Wrong vectors size')
assert(~isempty(vectors), 'Wrong vectors size')
% final array that contains all data
Nproj_orig = [cfg.iProjV, cfg.iProjU,cfg.iProjAngles];
Nvol_orig = size(volData);
%disp(split)
%disp(Nvol_orig)
Nvol_sub = Nvol_orig./split(1:3);
assert(all(mod(Nvol_sub,1)==0), 'Volume cannot be split')
% size of the subprojection of single subvolume
Nproj_sub = [ cfg.pixel_scale(1) * (Nvol_sub(3)* sind(cfg.lamino_angle) + sqrt(sum(Nvol_sub(1:2).^2))*cosd(cfg.lamino_angle)), ...
cfg.pixel_scale(1) * sqrt(sum(Nvol_sub(1:2).^2)) , ...
cfg.iProjAngles];
% adjust sub projection size to account for inplane rotation of the geometry
if cfg.tilt_angle ~= 0
Nproj_rot = [cosd(cfg.tilt_angle), -sind(cfg.tilt_angle); +sind(cfg.tilt_angle), cosd(cfg.tilt_angle)] * [0,0; Nproj_sub(1:2)];
% calculate projection window size after rotation
Nproj_rot = max(Nproj_rot) - min(Nproj_rot);
% add some extra padding
Nproj_sub(1:2) = 2*(Nproj_rot - Nproj_sub(1:2)) + Nproj_sub(1:2);
end
% provide extra space for subpixel (linear) interpolation at the borders of the split volumes, needed only for noninteger CoR_offset
Nproj_sub = Nproj_sub + 2*([split(3), max(split([1,2])),0] - 1);
Nproj_sub(1:2) = ceil(Nproj_sub(1:2)/16)*16; % make it easier splitable for ASTRA
% only if there is not split in the horizontal dimension
if all(split(1:2) == 1)
% do not take larger than size of the Nproj_orig because the edges
% are not need anyway
Nproj_sub(2) = min(Nproj_sub(2), Nproj_orig(2));
end
split(4) = ceil(split(4) * 4*prod(Nproj_sub)/(r.max_memory_blocks/N_GPU));
cfg_small = cfg;
cfg_small.iVolX = cfg.iVolX/split(1);
cfg_small.iVolY = cfg.iVolY/split(2);
cfg_small.iVolZ = cfg.iVolZ/split(3);
% get new size of projections
cfg_small.iProjU = Nproj_sub(2);
cfg_small.iProjV = Nproj_sub(1);
% calculate and store offset of the center of rotation, it will be used later
offset = vectors(:,4:6) +(vectors(:,10:12).*cfg.iProjV/2+vectors(:,7:9).*cfg.iProjU/2 );
CoR_offset = -[dot(offset', vectors(:,10:12)') ./ dot(vectors(:,10:12)', vectors(:,10:12)');
dot(offset', vectors(:,7:9)') ./ dot(vectors(:,7:9)', vectors(:,7:9)')]' ;
% find vector that will shift subvolume into center of the new projection size
shift_vec = vectors(:,10:12)*(cfg.iProjV/2-cfg_small.iProjV/2)+vectors(:,7:9)*(cfg.iProjU/2-cfg_small.iProjU/2);
% remove centering offset && apply new offset
vectors(:,4:6) = vectors(:,4:6) + shift_vec;
if length(cfg.pixel_scale) == 2 && cfg.pixel_scale(1) ~= cfg.pixel_scale(2)
error('Variable pixel size for each axis is not implemented')
end
Nblocks = prod(split);
% split into volume cubes
projData_all = zeros(Nproj_orig, 'single');
if N_GPU > 1
poolobj = gcp('nocreate');
if isempty(poolobj) || poolobj.NumWorkers < N_GPU
delete(poolobj);
poolobj = parpool(N_GPU);
end
poolobj.IdleTimeout = 600; % set idle timeout to 10 hours
end
% when the function is finished, make sure to execute following code
global status
status = true;
if r.use_shared_memory
out = onCleanup(@()myCleanupFun());
end
% run blocks in series
% run sub-blocks on each GPU in parallel
t_total = tic();
clear output
%% START OF OUTER GPU LOOP
outputs_blocks = [];
%% unitialize one solver per GPU
for thread_id = 1:N_GPU
% parse inputs and try to split them if possible
[outputs_blocks, inputs_block{thread_id},cfg_all{thread_id}] = ...
submit_block(thread_id, thread_id, outputs_blocks, volData, cfg, cfg_small, vectors,CoR_offset, split, r, varargin{:} );
end
unprocessed_blocks = N_GPU+1:Nblocks;
% write back to the shared array projData_all
%% merge blocks back from GPUs and write to the shared array volData
for ii = 1:Nblocks
if r.verbose>0; utils.progressbar(ii, Nblocks); end
% set values from the small blocks to the final output arrays
[thread_id, timing, id] = gather_block( outputs_blocks,projData_all, cfg_all);
if ~isempty(unprocessed_blocks)
block_id = unprocessed_blocks(1);
unprocessed_blocks(1) = [];
if isa(outputs_blocks, 'parallel.FevalFuture') && sum([outputs_blocks.Read]) ~= 1
outputs_blocks
keyboard
end
% submit a new job once the previous is finished
[outputs_blocks, inputs_block{thread_id},cfg_all{block_id}] = ...
submit_block(block_id, thread_id, outputs_blocks, volData, cfg, cfg_small, vectors,CoR_offset, split, r, varargin{:} );
if isa(outputs_blocks, 'parallel.FevalFuture') && any(cat(1,[outputs_blocks.Read]))
outputs_blocks
keyboard
end
end
end
t_total = toc(t_total);
if r.verbose > 1
fprintf('Timing system: GPU init %3.2gs shared_mem down %3.2gs upload on GPU %3.2gs tomo projection %3.2gs download from GPU %3.2gs shared_mem %3.2gs \n',sum(timing,2) )
if length(r.GPU) > 1
fprintf('Timing local: GPU init %3.2gs shared_mem down %3.2gs upload on GPU %3.2gs tomo projection %3.2gs download from GPU %3.2gs shared_mem %3.2gs \n ',sum(timing,2)/max(1,length(r.GPU)) )
fprintf('Total time %3.2fs, parfor overhead %3.2fs \n', t_total, t_total - sum(sum(timing,2)/max(1,length(r.GPU))) )
end
end
% everything was fine -> no cleaning needed
status = false;
end
function [outputs_blocks,inputs_block, cfg_out] = submit_block(block_id, thread_id, outputs_blocks, volData, cfg, cfg_small, vectors, CoR_offset,split, r, varargin )
% prepare blocks for asynchonous processing
inputs_block = prepare_block(block_id, volData, cfg, cfg_small, vectors,CoR_offset, split, r, varargin{:});
N_GPU = length(r.GPU);
if isempty(outputs_blocks); clear outputs_blocks; end
cfg_out = inputs_block{2};
try
%% process preloaded data
% no parallel toolbox
if N_GPU <= 1
[outputs_blocks{thread_id}.projData_small,outputs_blocks{thread_id}.timing, outputs_blocks{thread_id}.id]=...
run_partial_projector(inputs_block, block_id, 1,r.GPU, 0);
else
% run it asynchronously
if r.verbose > 3
ticBytes(gcp);
end
outputs_blocks(thread_id) = parfeval(@run_partial_projector, 3, inputs_block,block_id,thread_id, r.GPU, 0);
if r.verbose > 3
try; tocBytes(gcp); end
end
end
catch err
disp(getReport(err))
utils.check_available_memory
keyboard
end
end
function [thread_id, timing, id] = gather_block(output_package,projData_all, cfg_all )
if isa(output_package, 'parallel.FevalFuture')
%% my version of the fetchNext function, it seems faster
id = [];
assert(any(~[output_package.Read]), 'All blocks are already read')
while true
for thread_id =1:length(output_package)
if strcmpi(output_package(thread_id).State, 'finished') && output_package(thread_id).Read == 0
try
[projData_small,timing,id] = output_package(thread_id).fetchOutputs;
catch err
if strcmpi(err.identifier, 'parallel:fevalqueue:InvalidExecutionResult')
warning('Unknown error, trying to restart parpool')
delete(gcp('nocreate'));
end
if strcmpi(err.identifier, 'parallel:fevalqueue:InvalidExecutionResult')
warning('Unknown error, trying to restart parpool')
delete(gcp('nocreate'));
end
if ~isempty(output_package(thread_id).Diary)
fprintf('============ THREAD %i FAILED, OUTPUT: ============= \n', thread_id)
disp(output_package(thread_id).Diary)
end
fprintf('============ THREAD %i FAILED, ERROR: ============= \n', thread_id)
disp(getReport(output_package(thread_id).Error))
keyboard
output_package.cancel
rethrow(err)
end
break
end
end
if ~isempty(id); break; end
pause(0.01) % wait for the data to be prepared
end
elseif iscell(output_package)
thread_id = 1;
id = output_package{thread_id}.id;
projData_small = output_package{thread_id}.projData_small;
timing = output_package{thread_id}.timing;
else
disp('FAILED ?? ')
keyboard
end
if isempty(projData_small)
warning('ASTRA projection probably failed')
keyboard
end
% write back to the full array stored in RAM
if isa(projData_small, 'shm')
% load data from shared memory
[s, projData_small] = projData_small.attach;
s.protected = false; % release shared memory
elseif ~isnumeric(projData_small)
keyboard
end
% write the obtained projections back to the full projection array (projData_all)
utils.add_to_3D_projection(projData_small, projData_all, cfg_all{id}.projection_shift(cfg_all{id}.angle_ids,1:2), cfg_all{id}.angle_ids, true, false);
end
function myCleanupFun()
global status
if status
% destroy all shared memory that could have been left behind
!ipcs -m | cut -d' ' -f2 | grep '^[0-9]' | while read x; do ipcrm -m $x; done
end
end
function prepared_block = prepare_block(id, volData, cfg, cfg_small, vectors,CoR_offset, split,r, varargin)
Nblocks = prod(split(1:3));
Nvol_orig = size(volData);
Nvol_sub = Nvol_orig./split(1:3);
volData_small = zeros(Nvol_sub, 'single');
[i,j,k,angle_block_id] = ind2sub(split,id);
pos = [i,j,k];
% iterate over X,Y,Z axis
ind = cell(3,1);
shift = zeros(3,1);
for n = 1:3
ind{n} = max(1, 1+(pos(n)-1)*Nvol_sub(n)):min(pos(n)*Nvol_sub(n),Nvol_orig(n)); %% find optimal shift of the subvolume
if mod(split(n),2)==1 %% odd
shift(n) = (pos(n) - ceil(split(n)/2))*Nvol_sub(n);
else
shift(n) = (pos(n) - ceil(split(n)/2)-1/2)*Nvol_sub(n);
end
end
cfg_small.iVolX = length(ind{1});
cfg_small.iVolY = length(ind{2});
cfg_small.iVolZ = length(ind{3});
%% for splitting to angular blocks
assert(split(4) == 1 || isempty(r.deformation_fields), 'Deformation fields with angular splitting not supported')
Nangle_per_blocks = ceil(cfg.iProjAngles / split(4));
angle_ids = 1+(angle_block_id-1)*Nangle_per_blocks:min(cfg.iProjAngles, angle_block_id*Nangle_per_blocks);
cfg_small.iProjAngles = length(angle_ids);
cfg_small.angle_ids = angle_ids;
if ~isempty(r.deformation_fields)
for ii = 1:3
N_full = size(r.deformation_fields{1,ii});
N_small = ceil( N_full ./ reshape(split(1:3),[],1)');
for kk = 1:3
ind_def{kk} = (1+(pos(kk)-1)*N_small(kk)) : min(N_full(kk), (split(kk))*N_small(kk));
end
for jj = 1:2
deformation_fields_sub{jj,ii} = r.deformation_fields{jj,ii}(ind_def{:});
end
end
else
deformation_fields_sub = {};
end
CoR_offset = [CoR_offset, zeros(cfg.iProjAngles,1)] ;
vec = vectors;
% shift the sub-projections off center to create single
% large projection after assembling, if shift == 0 =>
% projections will be rotationally centered
% find optimal shift of the projections in the projData_all matrix
projection_shift = cfg_small.pixel_scale(1).^2* [vec(:,10:12)*shift, vec(:,7:9)*shift, zeros(cfg.iProjAngles,1)];
% change offset of the detector center
projection_shift = bsxfun(@plus, projection_shift , [cfg.iProjV/2-cfg_small.iProjV/2,cfg.iProjU/2-cfg_small.iProjU/2,0] + CoR_offset);
% calculate subpixel shifts
projection_shift_subpix = projection_shift - round(projection_shift) - CoR_offset;
projection_shift = round(projection_shift);
% apply subpixel shifts
vec(:,4:6) = vec(:,4:6) - ...
( bsxfun(@times,vectors(:,7:9),projection_shift_subpix(:,2))+ ...
bsxfun(@times,vectors(:,10:12),projection_shift_subpix(:,1)));
% just store it for later
cfg_small.projection_shift = projection_shift; % + [20,0,0];
% return only vector for the use angles
vec = vec(angle_ids,:);
prepared_block = {[], cfg_small, vec, r.split_sub, varargin{:},...
'verbose',0, 'deformation_fields', deformation_fields_sub, 'GPU', [], 'keep_on_GPU', false};
% take only small subvolume, (unfortunatelly this is more than duplicate the needed RAM !! )
if Nblocks > 1
% copy data from full volume into smaller field
% volData_small = volData(ind{:}); % move using matlab !! slow !!
% move data using custom made MEX routine
if r.use_shared_memory
s = shm();
s.allocate(volData_small);
% attach the shared memory
[s, volsmall_shm] = s.attach();
% === write data =====
% use self-made MEX OMP function to move the data
positions = ones(Nvol_sub(3),2,'int32').*int32([ind{1}(1),ind{2}(1)]-1);
% !! fill the data direclty to the shared memory
utils.get_from_3D_projection(volsmall_shm, volData,positions , int32(ind{3})');
% detach the shared memory
prepared_block{1} = s;
s.detach;
else
volData_small = volData(ind{:});
prepared_block{1} = volData_small;
end
else
prepared_block{1} = volData; % avoid memory copy if possible
end
end
function [projData_small, timing,block_id] = run_partial_projector(prepared_block, block_id,thread_id, GPU_list, verbose)
try
t0 = tic;
gpu = gpuDevice();
gpu_id = GPU_list(thread_id);
% let parfor to choose which GPU use
if gpu.Index ~= gpu_id
gpuDevice(gpu_id); % avoid unneeded initalization
end
t_init = toc(t0);
timing = [t_init, 0,0,0,0,0];
t = tic;
if isa(prepared_block{1}, 'shm')
% data are downloaded from shared memory
[s,volData_small] = prepared_block{1}.attach();
else
% data are given directly to the worker
volData_small = prepared_block{1};
end
timing(2) = toc(t);
t = tic;
timing(3) = toc(t);
t = tic;
is_remote = ~isempty(getCurrentTask());
% call the next level abstraction around ASTRA wrapper
projData_small = astra.Ax_partial(volData_small, prepared_block{2:end}, 'keep_on_GPU', true, 'verbose', is_remote);
timing(4) = toc(t);
t = tic;
projData_small = gather(projData_small); % move to RAM
timing(5) = toc(t);
t = tic;
if isa(prepared_block{1}, 'shm')
tic
s.detach;
% data are distributed to shared memory
s = shm(true);
s.upload(projData_small)
projData_small = s;
toc
end
timing(6) = toc(t);
catch err
gpu = gpuDevice
reset(gpu);
fprintf('Error on GPU %i / %i', gpu.Index, gpuDeviceCount)
disp( getReport(err, 'extended', 'hyperlinks', 'on'))
rethrow(err)
end
end