Files
2026-08-07 15:56:42 +09:00

553 lines
22 KiB
Matlab
Raw Permalink Blame History

This file contains invisible Unicode characters
This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
% 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