mirror of
https://github.com/c-sooyoung/fold_slice.git
synced 2026-09-17 23:39:08 +09:00
647 lines
29 KiB
Matlab
647 lines
29 KiB
Matlab
% BLOCK_FUN autosplitting function for fast GPU data processing.
|
||
% in the simplest case this function is equivalent to
|
||
% [varargout{:}] = fun(fp16.get(stack_object(ROI{:},:)), varargin{2:end});
|
||
% But for large datasets, it will split the stack_object along 3rd axis and try to also
|
||
% split other inputs if possible. Then it will submit each of the
|
||
% blocks on 1 or more GPUs and process. Processed data are returned back to RAM.
|
||
% it is useful even to process large blocks of data in RAM, to prevent too large RAM allocation.
|
||
% NOTE: Avoiding GPU processing ('use_gpu', false) can be useful for simple operations when the memory transfer to GPU is much slower than task execution !!
|
||
% NOTE: Multi GPU processing addes extra overhead for moving data to shared arrays and back, useful only for very expensive operations
|
||
% NOTE: If input is fp16 precision, it is automatically converted to singles for processing and result are converted back to fp16 when returned (unless ('use_fp16', false) is used)
|
||
% NOTE: Automatic splitting assumes memory required for two FFT operation and real to complex conversion with single precision numbers, ie MEMORY = 4*6*8*numel(stack_object)
|
||
%
|
||
% varargout = block_fun(fun, stack_object, varargin)
|
||
%
|
||
% Inputs:
|
||
% **fun - function to be executed
|
||
% **stack_object - 3D array that will be split along 3rd dimension and processed on GPU (as default)
|
||
% **varargin - other arguments, if their size along 1st dimension is the same as 3rd axis of stack_object, they will be split
|
||
% **parameter structure - if the last argument is structure,
|
||
% it will not be given to the "fun" but it will be used as a parameter input for block_fun.
|
||
% SEE DESCRIPTION IN CODE FOR ALL POSSIBLE INPUTS
|
||
% *returns*
|
||
% ++varargout - 3D array of size stack_object or 1D/2D with 1st dimension assumed to be equal to 3rd axis of stack_object
|
||
%
|
||
% Examples:
|
||
% % simple 1:1 array processing
|
||
% X = randn(Nx, Ny, Nz)
|
||
% X_fft = block_fun(@fft2,X); % equivalent to fft2(x) for small datasets
|
||
% % simple 1:1 array processing with manually defined blocks size , can be useful if autosplitting does not provide good results (ie too large / too small blocks)
|
||
% X_fft = block_fun(@fft2,X, struct('Nblocks', 10)); % equivalent to fft2(x) for small datasets
|
||
% % information reduction to (Nz,1,1),
|
||
% [x,y] = block_fun(@utils.center, X); % equivalent to utils.center(x) for small datasets
|
||
% % information reduction to (Nx, Nx,1)
|
||
% X_accum = block_fun(@(x)(abs(x).^2), X, struct('reduce_fun', 'plus')); % equivalent to sum(abs(x).^2,3) for small datasets
|
||
% % information reduction to (Nx, Nx,1) using only CPU
|
||
% X_accum = block_fun(@(x)(abs(x).^2), X, struct('use_gpu', false, 'reduce_fun', 'plus')); % equivalent to sum(abs(x).^2,3) for small datasets but avoids using GPU
|
||
% % information reduction to (Nx, Nx,1) and return results as single even if inputs is fp16
|
||
% X_accum = block_fun(@(x)(abs(x).^2), X, struct('use_fp16', false, 'reduce_fun', 'plus')); % equivalent to sum(abs(x).^2,3) for small datasets but avoids using GPU
|
||
% % use multiple inputs and perform inplace image shifting in order to prevent large memory allocation
|
||
% shift_x = randn(Nz,1); shift_y = randn(Nz,1);
|
||
% X_accum = block_fun(@utils.imshift_fft, X, shift_x, shift_y, struct('inplace', true)); % equivalent to imshift_fft(X, shift_x, shift_y)
|
||
% % process only a small subregion of input array, avoid slow matlab memory copy / allocation
|
||
% ROI = {1:10, 4:50};
|
||
% X_phase_cropped = block_fun( @angle, X, struct('ROI', ROI)); % angle(X(ROI{:}))
|
||
% % make the processing quiet, ie no progressbars are shown
|
||
% X_phase = block_fun( @angle, X, struct('verbose_level', 0)); % angle(X)
|
||
|
||
|
||
|
||
|
||
% clean all memory blocks:
|
||
% ! ipcs -m | cut -d' ' -f2 | grep '^[0-9]' | while read x; do ipcrm -m $x; done
|
||
|
||
|
||
%*-----------------------------------------------------------------------*
|
||
%| |
|
||
%| 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 varargout = block_fun(fun, stack_object, varargin)
|
||
|
||
global pprev
|
||
assert(~isempty(stack_object), 'Processed array is empty')
|
||
|
||
Nlayers = size(stack_object,3);
|
||
varargout = cell(nargout,1);
|
||
varargin = [{stack_object}, varargin];
|
||
% check if there is given parameter structure
|
||
if isstruct(varargin{end})
|
||
param = varargin{end};
|
||
varargin = varargin(1:end-1); % !!! remove the parameter structure from the vargin list
|
||
else
|
||
param = struct();
|
||
end
|
||
ninputs = length(varargin);
|
||
assert(~isempty(stack_object), 'Input array is empty')
|
||
|
||
%% ============= LOAD PARAMETERS FROM PARAM STRUCTURE OR USE PREDEFINED DEFAULTS ================
|
||
if ~isfield(param, 'verbose_level')
|
||
param.verbose_level = 1;
|
||
end
|
||
if ~isfield(param, 'use_GPU')
|
||
param.use_GPU = true; % some operations are just faster without moving to GPU
|
||
end
|
||
if ~isfield(param, 'move_to_GPU')
|
||
param.move_to_GPU = true; % move blocks directly yo GPU or let the called function to decide
|
||
end
|
||
if ~isfield(param, 'GPU_list') || isempty(param.GPU_list)
|
||
if param.use_GPU && gpuDeviceCount > 0
|
||
gpu = gpuDevice;
|
||
param.GPU_list = gpu.Index;
|
||
else
|
||
param.GPU_list = -1; % list of the used GPUs
|
||
end
|
||
end
|
||
if ~isfield(param, 'inplace')
|
||
param.inplace = false; % run the function inplace , !! dont stop the operation in middle of process !!
|
||
% assume that output is the same as first
|
||
% input array and it will write results
|
||
% into this memory without reallocating
|
||
end
|
||
|
||
if ~isfield(param, 'ROI') || isempty(param.ROI) % process only limited ROI of all inputs of size stack_object
|
||
param.ROI = {':', ':'};
|
||
end
|
||
Nelements = 1;
|
||
for i = 1:2
|
||
if strcmpi(param.ROI{i},':')
|
||
param.ROI{i} = 1:size(stack_object,i);
|
||
end
|
||
Nelements = Nelements * length(param.ROI{i});
|
||
end
|
||
assert(Nelements>0, 'Selected ROI is empty')
|
||
|
||
if ~isfield(param, 'Nblocks') % number of blocks to split the input volume, [] = auto
|
||
param.Nblocks = [];
|
||
end
|
||
if ~isfield(param, 'full_block_size') % largest volume size in the calculation, default = size(stack_object)
|
||
param.full_block_size = size(stack_object);
|
||
end
|
||
if ~isfield(param, 'use_fp16')
|
||
param.use_fp16 = isa(stack_object, 'uint16');
|
||
end
|
||
% function applied on the computed data to reduce along 3rd axis
|
||
if ~isfield(param, 'reduce_fun')
|
||
param.reduce_fun = [];
|
||
else
|
||
assert(ismember(func2str(param.reduce_fun), {'min','max','plus'}))
|
||
end
|
||
if ~isfield(param, 'use_shared_memory')
|
||
param.use_shared_memory = false; % use shared memory for data exchange, it is used only if N_GPU > 1 or for debugging
|
||
end
|
||
if param.inplace
|
||
varargout{1} = stack_object;
|
||
end
|
||
|
||
if param.inplace
|
||
warning on
|
||
warning off backtrace
|
||
warning('Inplace data processing, DO NOT INTERRUPT')
|
||
warning on backtrace
|
||
end
|
||
|
||
if (gpuDeviceCount == 0 || isa(stack_object, 'gpuArray') || ismatrix(stack_object)) && ...
|
||
(isempty(param.Nblocks) || param.Nblocks == 1)
|
||
%% ================ EXECUTE FUN =============================
|
||
[varargout{:}] = fun(stack_object(param.ROI{:},:), varargin{2:end});
|
||
else
|
||
%% GPU or CPU splitting is required
|
||
N_GPU = max(1,length(param.GPU_list));
|
||
if param.use_GPU
|
||
gpu = gpuDevice;
|
||
if ~ismember(gpu.Index,param.GPU_list)
|
||
gpu = gpuDevice(param.GPU_list(1));
|
||
end
|
||
end
|
||
if isempty(param.Nblocks)
|
||
%% autoestimate block size
|
||
if param.use_GPU
|
||
%% empirical condition assuming FFT involved, may be too pesimistic
|
||
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
||
Nblocks = ceil( (4*6*8* prod(param.full_block_size)) / gpu.AvailableMemory) ;
|
||
Nblocks = max(Nblocks, prod(param.full_block_size)/ double(intmax('int32')));
|
||
Nblocks = max(Nblocks, N_GPU);
|
||
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
||
else
|
||
%% if CPU is used, split it to <20GB blocks
|
||
max_block_size = min(utils.check_available_memory*1e6, 20e9); %% work with 10GB blocks
|
||
Nblocks = ceil( (6*8* prod(param.full_block_size)) / max_block_size) ;
|
||
end
|
||
else
|
||
Nblocks = param.Nblocks;
|
||
end
|
||
|
||
% optimize size of the block to equally use all provided GPUs
|
||
N = ceil(Nlayers / Nblocks);
|
||
Nblocks = ceil(Nlayers / N / N_GPU) * N_GPU; % make it splitable to N_GPU
|
||
N = ceil(Nlayers / Nblocks);
|
||
|
||
% parse inputs and store infomations needed to split the inputs
|
||
splitable_along_3rd_axis = false(ninputs,1);
|
||
apply_ROI = false(ninputs,1);
|
||
|
||
for jj = 1:ninputs
|
||
splitable_along_3rd_axis(jj) = ndims(varargin{jj}) == 3 && (size(varargin{jj},3) == size(stack_object,3));
|
||
apply_ROI(jj) = all([size(varargin{jj},1),size(varargin{jj},2)] == [size(stack_object,1), size(stack_object,2)]) && ...
|
||
((size(varargin{jj},1) ~= length(param.ROI{1})) || (size(varargin{jj},2) ~= length(param.ROI{2})));
|
||
end
|
||
|
||
% expected inputs are singles or uint16, if they are doubles, convert to single so save memory and speed up execution on GPU
|
||
for jj = find(splitable_along_3rd_axis)'
|
||
if ~ismember(class(varargin{jj}), {'single', 'uint16', 'uint8'})
|
||
varargin{jj} = single(varargin{jj});
|
||
end
|
||
end
|
||
|
||
if Nblocks == 1
|
||
%% in case no splitting is needed
|
||
varargout = process_in_single_block(fun, varargin, apply_ROI,splitable_along_3rd_axis, nargout, param);
|
||
return
|
||
|
||
else
|
||
|
||
% when the function is finished, make sure to execute following code
|
||
global status
|
||
status = true;
|
||
|
||
% avoid using shared memory if not needed, it is a bit slower
|
||
use_shared_memory = length(param.GPU_list) > 1 || param.use_shared_memory;
|
||
if use_shared_memory
|
||
out = onCleanup(@myCleanupFun);
|
||
end
|
||
for ii = 1:length(varargin)
|
||
if splitable_along_3rd_axis(ii) || N_GPU > 1
|
||
varargin{ii} = gather(varargin{ii});
|
||
end
|
||
end
|
||
|
||
ind = {};
|
||
|
||
for ii = 1:Nblocks
|
||
ind{end+1} = 1+N*(ii-1) : min(Nlayers, N*ii);
|
||
% avoid single layer arrays
|
||
if length(ind{end}) < 2
|
||
ind{end-1} = [ind{end-1}, ind{end}];
|
||
ind(end) = [];
|
||
end
|
||
end
|
||
Nblocks = length(ind);
|
||
|
||
%% open parpool to allow multi GPU processing
|
||
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
|
||
|
||
if param.verbose_level>0
|
||
pprev = -1;
|
||
end
|
||
%% START OF OUTER GPU LOOP
|
||
outputs_blocks = [];
|
||
%% unitialize one solver per GPU
|
||
for block_id = 1:N_GPU
|
||
% parse inputs and try to split them if possible
|
||
[outputs_blocks, inputs_block{block_id}] = submit_block(block_id,block_id, ind, varargin,outputs_blocks ,param, fun, apply_ROI, use_shared_memory, Nlayers,nargout, splitable_along_3rd_axis );
|
||
end
|
||
|
||
unprocessed_blocks = N_GPU+1:Nblocks;
|
||
|
||
%% merge blocks back from GPUs
|
||
for ii = 1:Nblocks
|
||
if param.verbose_level>0; utils.progressbar(ii, Nblocks); end
|
||
% set values from the small blocks to the final output arrays
|
||
[thread_id, varargout] = gather_block(varargout,param,Nlayers, nargout,ind, N, outputs_blocks);
|
||
if ~isempty(unprocessed_blocks)
|
||
block_id = unprocessed_blocks(1);
|
||
unprocessed_blocks(1) = [];
|
||
% submit a new job once the previous is finished
|
||
[outputs_blocks, inputs_block{thread_id}] = submit_block(block_id,thread_id,ind, varargin,outputs_blocks ,param, fun, apply_ROI, use_shared_memory, Nlayers,nargout, splitable_along_3rd_axis );
|
||
if isa(outputs_blocks, 'parallel.FevalFuture') && any(cat(1,[outputs_blocks.Read]))
|
||
outputs_blocks
|
||
warning on
|
||
warning('Unknown error in parallel processing')
|
||
keyboard
|
||
end
|
||
end
|
||
|
||
end
|
||
|
||
% END OF OUTER GPU LOOP
|
||
end
|
||
end
|
||
|
||
% everything was fine -> no cleaning needed
|
||
status = false;
|
||
|
||
if param.inplace
|
||
warning on
|
||
warning off backtrace
|
||
warning('Inplace data processing finished')
|
||
warning on backtrace
|
||
end
|
||
end
|
||
|
||
%% auxiliary function
|
||
|
||
function [outputs_blocks,inputs_block] = submit_block(block_id, thread_id, ind, inputs, outputs_blocks, param, fun, apply_ROI, use_shared_memory, Nlayers , Noutputs, splitable_along_3rd_axis )
|
||
inputs_block = inputs;
|
||
ind = ind{block_id};
|
||
% get blocks of data
|
||
for jj = 1:length(inputs)
|
||
if splitable_along_3rd_axis(jj)
|
||
if apply_ROI(jj)
|
||
ROI = param.ROI;
|
||
else
|
||
ROI = {};
|
||
end
|
||
% use fast splitting based on MEX files for large splitable blocks
|
||
inputs_block{jj} = get_block(inputs_block{jj}, ind, ROI,use_shared_memory);
|
||
elseif isnumeric(inputs_block{jj}) && size(inputs_block{jj},1) == Nlayers
|
||
% !! assume split along first dimension of the object !!!
|
||
inputs_block{jj} = inputs_block{jj}(ind,:,:,:); % split it and provide subfunction only the valid chunk
|
||
end
|
||
end
|
||
if isempty(outputs_blocks); clear outputs_blocks; end
|
||
%% ----------- call the function "fun" , execute on GPUs -----------------
|
||
if length(param.GPU_list) <= 1 || ~param.use_GPU
|
||
%% single GPU
|
||
[outputs_blocks.data,outputs_blocks.id]...
|
||
= worker(fun,Noutputs, use_shared_memory,param,block_id,1,inputs_block);
|
||
else
|
||
%% run asynchronous on multiple GPUs
|
||
outputs_blocks(thread_id) = parfeval(@worker,2,fun,Noutputs, use_shared_memory,param,block_id,thread_id,inputs_block);
|
||
end
|
||
end
|
||
|
||
function out = get_block(full_array, ind, ROI, use_shared_memory)
|
||
Nlayers = length(ind);
|
||
positions = zeros(Nlayers,2,'int32');
|
||
if ~isempty(ROI)
|
||
Npix_small = [length(ROI{1}),length(ROI{2}),length(ind)];
|
||
assert(all(Npix_small <= size(full_array)), 'Provided ROI is larger than input array')
|
||
positions = positions+int32([ROI{1}(1),ROI{2}(1)]-1);
|
||
else
|
||
Npix_small = [size(full_array,1),size(full_array,2),length(ind)];
|
||
end
|
||
% allocate in RAM array to load the block from "stack_object"
|
||
stack_object_block_in = zeros(Npix_small, 'like', full_array);
|
||
|
||
if use_shared_memory
|
||
s = shm(); % create sharemem object
|
||
if isa(full_array, 'single')
|
||
s.allocate(stack_object_block_in); % allocate SHM memory
|
||
% attach the shared memory
|
||
[s, stack_object_block_shm] = s.attach();
|
||
% === write data =====
|
||
% use self-made MEX OMP function to move the data
|
||
utils.get_from_3D_projection(stack_object_block_shm, full_array, positions, ind);
|
||
elseif isa(full_array, 'uint16')
|
||
stack_object_block_in = zeros(Npix_small, 'like', full_array);
|
||
% === write data =====
|
||
% use self-made MEX OMP function to move the data
|
||
utils.get_from_3D_projection(stack_object_block_in, full_array, positions, ind);
|
||
% convert from fp16 to singles
|
||
stack_object_block_in = fp16.get(stack_object_block_in);
|
||
|
||
% upload to the shared memory
|
||
s.upload(stack_object_block_in); % seems to be really slow ????
|
||
|
||
%s.allocate(stack_object_block_in); % allocate SHM memory
|
||
%[s, stack_object_block_shm] = s.attach();
|
||
% % move to shared memory , a bit faster than the memcpy in the
|
||
% % sharedmem function
|
||
%utils.get_from_3D_projection(stack_object_block_shm, stack_object_block_in, zeros(Nlayers,2,'int32'), int32(1:Nlayers));
|
||
%toc
|
||
end
|
||
% detach the shared memory
|
||
s.detach;
|
||
% return structure with shared mem object
|
||
out = s;
|
||
else
|
||
% use custom made MEX OMP function to move the data
|
||
stack_object_block_in = utils.get_from_3D_projection(stack_object_block_in, full_array, positions, ind);
|
||
% if needed, convert from fp16 to singles
|
||
out = fp16.get(stack_object_block_in);
|
||
end
|
||
end
|
||
|
||
function myCleanupFun()
|
||
global status
|
||
if status
|
||
warning('!! cleaning all shared memory !!')
|
||
% 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 [thread_id, outputs] = gather_block(outputs,param,Nlayers,Noutputs, ind_all, Nl_per_block, output_package)
|
||
if isa(output_package, 'parallel.FevalFuture')
|
||
% gather results from cluster , WAIT FOR CALCULATIONS TO BE FINISHED
|
||
% [~, outputs_block, id] = fetchNext(output_package);
|
||
%% my version of the fetchNext function, it seems faster
|
||
id = [];
|
||
while true
|
||
for thread_id =1:length(output_package)
|
||
if strcmpi(output_package(thread_id).State, 'finished') && output_package(thread_id).Read == 0
|
||
try
|
||
[outputs_block, 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 ~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
|
||
else
|
||
thread_id = 1;
|
||
id = output_package.id;
|
||
outputs_block = output_package.data;
|
||
end
|
||
|
||
ind = ind_all{id};
|
||
|
||
|
||
for jj = 1:Noutputs
|
||
% get the results from GPU to RAM
|
||
% allocate RAM for the final output arrays , allocate when
|
||
% with the same complexity as the outputs_block{ii}{jj} has
|
||
s = shm();
|
||
if isa(outputs_block{jj}, 'shm')
|
||
% first download from the shared mem back to workspace
|
||
[s, outputs_block{jj}] = outputs_block{jj}.attach;
|
||
s.protected = false;
|
||
end
|
||
|
||
% processing returns the same number of layers as input
|
||
is_block = ndims(outputs_block{jj}) == 3 && size(outputs_block{jj},3) == length(ind) ;
|
||
% processing returns one layer per input
|
||
is_reduced = ismatrix(outputs_block{jj}) && size(outputs_block{jj},3) == 1 && size(outputs_block{jj},1) ~= length(ind) && ~isempty(param.reduce_fun) && jj ==1;
|
||
if isa(outputs_block{jj}, 'double')
|
||
% small arrays / scalars just convert
|
||
outputs_block{jj} = single(outputs_block{jj});
|
||
% for large complain
|
||
if is_block
|
||
warning('Output of the function should be blocks in single precision')
|
||
end
|
||
end
|
||
if is_block && param.use_fp16 && isa(outputs_block{jj},'single')
|
||
% move to fp16 precision to avoid large memory allocation
|
||
outputs_block{jj} = fp16.set(outputs_block{jj});
|
||
end
|
||
|
||
|
||
% if the first iteration when it is empty create the large output arrays
|
||
if isempty(outputs{jj})
|
||
try
|
||
if is_block
|
||
% return large array of size stack_object
|
||
Npix_small = [size(outputs_block{jj},1),size(outputs_block{jj},2),Nlayers];
|
||
outputs{jj} = zeros(Npix_small, 'like', outputs_block{jj});
|
||
elseif is_reduced
|
||
outputs{jj} = [];
|
||
elseif size(outputs_block{jj},1) == length(ind) && ndims(outputs_block{jj}) <= 3
|
||
% return smaller array split along first dimension
|
||
outputs{jj} = zeros([Nlayers, size(outputs_block{jj},2), size(outputs_block{jj},3)], 'like', outputs_block{jj});
|
||
else
|
||
error('Unimplemented option')
|
||
end
|
||
catch err
|
||
utils.check_available_memory
|
||
rethrow(err)
|
||
end
|
||
end
|
||
|
||
if is_block
|
||
% if the output is 3D array, add the results into the large output stored in RAM
|
||
outputs{jj} = tomo.set_to_array(outputs{jj}, outputs_block{jj}, Nl_per_block*(id-1), false);
|
||
elseif is_reduced
|
||
if isempty(outputs{jj})
|
||
outputs{jj} = outputs_block{jj};
|
||
else
|
||
outputs{jj} = param.reduce_fun(outputs{jj}, outputs_block{jj});
|
||
end
|
||
else
|
||
outputs{jj}(ind,:,:) = outputs_block{jj};
|
||
end
|
||
end
|
||
end
|
||
|
||
function [outputs,id] = worker(fun, Noutputs, use_sharemem, param, id, thread_id, inputs)
|
||
|
||
% let parfor to choose which GPU use
|
||
gpu_id = param.GPU_list(thread_id);
|
||
gpu = gpuDevice();
|
||
if ~isempty(param.GPU_list) && gpu.Index ~= gpu_id && gpu_id > 0
|
||
gpuDevice(gpu_id); % avoid unneeded initalization
|
||
end
|
||
|
||
if ~isempty(getCurrentTask())
|
||
% report only if inside parfor
|
||
utils.check_available_memory
|
||
if gpu_id > 0
|
||
fprintf('GPU %i ===== Memory %3.2GB/%3.2gGB \n',gpu.Index, gpu.AvailableMemory/1e9, gpu.TotalMemory/1e9 )
|
||
end
|
||
end
|
||
|
||
%% upload data on GPU
|
||
|
||
for jj = 1:length(inputs)
|
||
if isa(inputs{jj}, 'shm')
|
||
% data are downloaded from shared memory
|
||
[s,inputs{jj}]=inputs{jj}.attach;
|
||
end
|
||
|
||
if isnumeric(inputs{jj}) && numel(inputs{jj}) > 1e4 && param.use_GPU && param.move_to_GPU
|
||
inputs{jj} = gpuArray(inputs{jj});
|
||
end
|
||
end
|
||
outputs = cell(Noutputs,1);
|
||
Nlayers = size(inputs{1},3);
|
||
try
|
||
%%%%%% CALL THE WORKING FUNCTION WITH ARRAYS ALREADY MOVED TO GPU %%%%
|
||
[outputs{:}] = fun(inputs{:});
|
||
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
|
||
|
||
if ~isempty(param.reduce_fun)
|
||
% apply additional function to reduce the data
|
||
switch func2str(param.reduce_fun)
|
||
case 'plus',outputs{1} = sum(outputs{1},3);
|
||
case 'min', outputs{1} = min(outputs{1},[],3);
|
||
case 'max', outputs{1} = max(outputs{1},[],3);
|
||
end
|
||
end
|
||
catch err
|
||
|
||
if any(strcmpi(err.identifier, {'parallel:gpu:array:OOMForOperation', 'parallel:gpu:array:OOM'}))
|
||
[~,out] = system(sprintf('nvidia-smi --id=%i | head -n10 | tail -n7', gpu.Index-1));
|
||
% disp(getReport(err))
|
||
|
||
disp(out)
|
||
warning('Low memory on GPU %i, RESETTING GPU ... \n', gpu.Index)
|
||
|
||
for jj = 1:length(inputs)
|
||
inputs{jj} = gather(inputs{jj});
|
||
end
|
||
% try to clean some leftover data from previous calculations
|
||
reset(gpuDevice)
|
||
|
||
fprintf('TRYING TO RECURSIVELLY CALL BLOCKFUN %s WITH 2 SUBBLOCKS \n', func2str(fun))
|
||
[outputs{:}] = tomo.block_fun(fun,inputs{:},struct('Nblocks',2));
|
||
else
|
||
rethrow(err)
|
||
end
|
||
end
|
||
|
||
|
||
for jj = 1:length(outputs)
|
||
%% return data from GPU
|
||
if isa(outputs{jj}, 'gpuArray') && param.use_GPU
|
||
outputs{jj} = gather(outputs{jj});
|
||
end
|
||
%% move data to shared memory if needed
|
||
is_block = ndims(outputs{jj}) == 3 && (size(outputs{jj},3) == Nlayers);
|
||
if is_block && use_sharemem
|
||
% data are distributed to shared memory
|
||
s = shm(true);
|
||
s.upload(outputs{jj})
|
||
outputs{jj} = s;
|
||
end
|
||
end
|
||
end
|
||
|
||
function outputs = process_in_single_block(fun, inputs,apply_ROI, splitable_along_3rd_axis, Noutputs, param)
|
||
for jj = 1:length(inputs)
|
||
if splitable_along_3rd_axis(jj)
|
||
% if stored as fp16 precision, convert to singles
|
||
inputs{jj} = fp16.get( inputs{jj});
|
||
% if provided and possible, use param.ROI to crop the working / inputs array
|
||
if apply_ROI(jj)
|
||
inputs{jj} = inputs{jj}(param.ROI{:},:);
|
||
end
|
||
end
|
||
if isa(inputs{jj}, 'double')
|
||
inputs{jj} = single(inputs{jj});
|
||
end
|
||
|
||
if param.use_GPU && isnumeric(inputs{jj}) && numel(inputs{jj}) > 1e2 && param.move_to_GPU
|
||
% upload on GPU
|
||
inputs{jj} = gpuArray(inputs{jj});
|
||
end
|
||
end
|
||
outputs = cell(Noutputs,1);
|
||
% execute the worker
|
||
[outputs{:}] = fun(inputs{:});
|
||
if ~isempty(param.reduce_fun)
|
||
% apply additional function to reduce the data
|
||
switch func2str(param.reduce_fun)
|
||
case 'plus',outputs{1} = sum(outputs{1},3);
|
||
case 'min', outputs{1} = min(outputs{1},[],3);
|
||
case 'max', outputs{1} = max(outputs{1},[],3);
|
||
end
|
||
end
|
||
for jj = 1:Noutputs
|
||
if param.use_GPU && isa(outputs{jj}, 'gpuArray')
|
||
% get back from GPU
|
||
outputs{jj} = gather(outputs{jj});
|
||
end
|
||
% processing returns the same number of layers as input
|
||
is_block = ndims(outputs{jj}) == 3 && size(outputs{jj},3) == size(inputs{1},3);
|
||
if param.use_fp16 && is_block
|
||
outputs{jj} = fp16.set(outputs{jj});
|
||
end
|
||
end
|
||
end
|