diff --git a/idg-api/Buffer.cpp b/idg-api/Buffer.cpp index b2628506eb57e25337f34ff6a6060b88197b2ab2..9f7a31be74f694e861a76169c1edc42488828222 100644 --- a/idg-api/Buffer.cpp +++ b/idg-api/Buffer.cpp @@ -31,14 +31,16 @@ BufferImpl::BufferImpl(const BufferSetImpl& bufferset, size_t bufferTimesteps) m_timeStartNextBatch(bufferTimesteps), m_nrStations(0), m_nr_baselines(0), - m_shift(2), + m_shift{0, 0}, m_default_aterm_offsets(2), - m_aterm_offsets_array(0), - m_frequencies(0), - m_aterms_array(0, 0, 0, 0), - m_bufferUVW(0, 0), - m_bufferStationPairs(0), - m_bufferVisibilities(0, 0, 0) { + m_frequencies(aocommon::xt::CreateSpan<float, 1>(nullptr, {0})), + m_bufferUVW( + aocommon::xt::CreateSpan<idg::UVW<float>, 2>(nullptr, {0, 0})), + m_bufferStationPairs( + aocommon::xt::CreateSpan<std::pair<unsigned int, unsigned int>, 1>( + nullptr, {0})), + m_bufferVisibilities(aocommon::xt::CreateSpan<std::complex<float>, 4>( + nullptr, {0, 0, 0, 0})) { #if defined(DEBUG) cout << __func__ << endl; #endif @@ -77,7 +79,7 @@ void BufferImpl::set_frequencies(size_t nr_channels, const double* frequencyList) { m_nr_channels = nr_channels; m_frequencies = - m_bufferset.get_proxy().allocate_array1d<float>(m_nr_channels); + m_bufferset.get_proxy().allocate_span<float, 1>({m_nr_channels}); for (int i = 0; i < m_nr_channels; i++) { m_frequencies(i) = frequencyList[i]; } @@ -86,19 +88,15 @@ void BufferImpl::set_frequencies(size_t nr_channels, void BufferImpl::set_frequencies(const std::vector<double>& frequency_list) { m_nr_channels = frequency_list.size(); m_frequencies = - m_bufferset.get_proxy().allocate_array1d<float>(m_nr_channels); - for (int i = 0; i < m_nr_channels; i++) { - m_frequencies(i) = frequency_list[i]; - } + m_bufferset.get_proxy().allocate_span<float, 1>({m_nr_channels}); + m_frequencies = xt::adapt(frequency_list); } double BufferImpl::get_frequency(const size_t channel) const { return m_frequencies(channel); } -size_t BufferImpl::get_frequencies_size() const { - return m_frequencies.get_x_dim(); -} +size_t BufferImpl::get_frequencies_size() const { return m_frequencies.size(); } // Plan creation and helper functions @@ -113,31 +111,30 @@ void BufferImpl::bake() { } void BufferImpl::malloc_buffers() { - const int nr_correlations = m_bufferset.get_nr_correlations(); + const size_t nr_correlations = m_bufferset.get_nr_correlations(); proxy::Proxy& proxy = m_bufferset.get_proxy(); m_bufferUVW = - proxy.allocate_array2d<UVW<float>>(m_nr_baselines, m_bufferTimesteps); - m_bufferVisibilities = proxy.allocate_array4d<std::complex<float>>( - m_nr_baselines, m_bufferTimesteps, m_nr_channels, nr_correlations); + proxy.allocate_span<UVW<float>, 2>({m_nr_baselines, m_bufferTimesteps}); + m_bufferVisibilities = proxy.allocate_span<std::complex<float>, 4>( + {m_nr_baselines, m_bufferTimesteps, m_nr_channels, nr_correlations}); m_bufferStationPairs = - proxy.allocate_array1d<std::pair<unsigned int, unsigned int>>( - m_nr_baselines); - m_bufferStationPairs.init({m_nrStations, m_nrStations}); - // already done: m_spheroidal.reserve(subgridsize, subgridsize); - // m_aterms = Array4D<Matrix2x2<std::complex<float>>>(1, m_nrStations, - // subgridsize, subgridsize); + proxy.allocate_span<std::pair<unsigned int, unsigned int>, 1>( + {m_nr_baselines}); + m_bufferStationPairs.fill( + std::pair<unsigned int, unsigned int>(m_nrStations, m_nrStations)); + // m_aterms is already allocated in BufferImpl::init_default_aterm } void BufferImpl::reset_buffers() { - m_bufferVisibilities.zero(); + m_bufferVisibilities.fill(std::complex<float>(0, 0)); set_uvw_to_infinity(); init_default_aterm(); } void BufferImpl::set_uvw_to_infinity() { - m_bufferUVW.init({std::numeric_limits<float>::infinity(), - std::numeric_limits<float>::infinity(), - std::numeric_limits<float>::infinity()}); + m_bufferUVW.fill(UVW<float>{std::numeric_limits<float>::infinity(), + std::numeric_limits<float>::infinity(), + std::numeric_limits<float>::infinity()}); } void BufferImpl::init_default_aterm() { diff --git a/idg-api/BufferImpl.h b/idg-api/BufferImpl.h index 9d8a237d51b2ff3fb483a0d57fa3b46eb8cca189..8cddf153b8f79e564aae78088ba31e40d0557b02 100644 --- a/idg-api/BufferImpl.h +++ b/idg-api/BufferImpl.h @@ -99,20 +99,20 @@ class BufferImpl : public virtual Buffer { size_t m_nrStations; size_t m_nr_channels; size_t m_nr_baselines; - Array1D<float> m_shift; + std::array<float, 2> m_shift; std::vector<unsigned int> m_default_aterm_offsets; std::vector<unsigned int> m_aterm_offsets; - Array1D<unsigned int> m_aterm_offsets_array; // Buffers - Array1D<float> m_frequencies; // CH + aocommon::xt::Span<float, 1> m_frequencies; // CH std::vector<Matrix2x2<std::complex<float>>> m_aterms; std::vector<Matrix2x2<std::complex<float>>> m_default_aterms; - Array4D<Matrix2x2<std::complex<float>>> m_aterms_array; // ST x SB x SB - Array2D<UVW<float>> m_bufferUVW; // BL x TI - Array1D<std::pair<unsigned int, unsigned int>> m_bufferStationPairs; // BL - Array4D<std::complex<float>> m_bufferVisibilities; // BL x TI x CH x CR + aocommon::xt::Span<UVW<float>, 2> m_bufferUVW; // BL x TI + aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1> + m_bufferStationPairs; // BL + aocommon::xt::Span<std::complex<float>, 4> + m_bufferVisibilities; // BL x TI x CH x CR }; } // namespace api diff --git a/idg-api/BufferSet.cpp b/idg-api/BufferSet.cpp index da15dd8a80e6ca11242ed8bc348d2ebdc78fe9e8..d86524bafac64b3b9c9961b521c58a4fd7c2f1ab 100644 --- a/idg-api/BufferSet.cpp +++ b/idg-api/BufferSet.cpp @@ -70,8 +70,7 @@ int nextcomposite(int n) { } BufferSetImpl::BufferSetImpl(Type architecture) - : m_default_aterm_correction(0, 0, 0, 0), - m_avg_aterm_correction(0, 0, 0, 0), + : m_spheroidal(aocommon::xt::CreateSpan<float, 2>(nullptr, {0, 0})), m_stokes_I_only(false), m_nr_correlations(4), m_nr_polarizations(4), @@ -91,7 +90,6 @@ BufferSetImpl::~BufferSetImpl() { m_gridderbuffers.clear(); m_degridderbuffers.clear(); m_bulkdegridders.clear(); - m_spheroidal.free(); m_proxy.reset(); report_runtime(); } @@ -269,17 +267,6 @@ void BufferSetImpl::init(size_t size, float cell_size, float max_w, int(std::ceil((m_kernel_size + uv_span_time + uv_span_frequency) / 8.0)) * 8; - m_default_aterm_correction = - Array4D<std::complex<float>>(m_subgridsize, m_subgridsize, 4, 4); - m_default_aterm_correction.init(0.0); - for (size_t i = 0; i < m_subgridsize; i++) { - for (size_t j = 0; j < m_subgridsize; j++) { - for (size_t k = 0; k < 4; k++) { - m_default_aterm_correction(i, j, k, k) = 1.0; - } - } - } - allocate_grid(); m_proxy->init_cache(m_subgridsize, m_cell_size, m_w_step, m_shift); @@ -306,7 +293,8 @@ void BufferSetImpl::init(size_t size, float cell_size, float max_w, } // Generate spheroidal using m_taper_subgrid. - m_spheroidal = m_proxy->allocate_array2d<float>(m_subgridsize, m_subgridsize); + m_spheroidal = + m_proxy->allocate_span<float, 2>({m_subgridsize, m_subgridsize}); for (size_t y = 0; y < m_subgridsize; y++) { for (size_t x = 0; x < m_subgridsize; x++) { m_spheroidal(y, x) = m_taper_subgrid[y] * m_taper_subgrid[x]; @@ -421,14 +409,15 @@ void BufferSetImpl::set_image(const double* image, bool do_scale) { arr_float_1D_t& phasor_imag __attribute__((aligned(64))) = *reinterpret_cast<arr_float_1D_t*>(aligned_alloc(64, size_1D)); + auto image_array = aocommon::xt::CreateSpan( + image, std::array<size_t, 3>{static_cast<size_t>(m_nr_polarizations), + m_size, m_size}); + #pragma omp for for (int y = 0; y < m_size; y++) { memset(w0_row_real, 0, m_nr_polarizations * m_size * sizeof(float)); memset(w0_row_imag, 0, m_nr_polarizations * m_size * sizeof(float)); - const Array3D<double> image_array(const_cast<double*>(image), - m_nr_polarizations, m_size, m_size); - // Copy row of image and convert stokes to polarizations if (m_stokes_I_only) { for (int x = 0; x < m_size; x++) { @@ -652,11 +641,12 @@ void BufferSetImpl::get_image(double* image) { arr_float_1D_t& phasor_imag __attribute__((aligned(64))) = *reinterpret_cast<arr_float_1D_t*>(aligned_alloc(64, size_1D)); + auto image_array = aocommon::xt::CreateSpan( + image, std::array<size_t, 3>{static_cast<size_t>(m_nr_polarizations), + m_size, m_size}); + #pragma omp for for (int y = 0; y < m_size; y++) { - Array3D<double> image_array((double*)image, m_nr_polarizations, m_size, - m_size); - // Compute inverse spheroidal for (int x = 0; x < m_size; x++) { inv_tapers[x] = m_inv_taper[y] * m_inv_taper[x]; @@ -944,9 +934,10 @@ void BufferSetImpl::finalize_compute_avg_beam() { } #endif - m_avg_aterm_correction = Array4D<std::complex<float>>( - m_matrix_inverse_beam->data(), m_subgridsize, m_subgridsize, 4, 4); - m_proxy->set_avg_aterm_correction(m_avg_aterm_correction); + auto matrix_inverse_beam_span = + aocommon::xt::CreateSpan<std::complex<float>, 4>( + m_matrix_inverse_beam->data(), {m_subgridsize, m_subgridsize, 4, 4}); + m_proxy->set_avg_aterm_correction(matrix_inverse_beam_span); #ifndef NDEBUG { @@ -962,14 +953,14 @@ void BufferSetImpl::finalize_compute_avg_beam() { void BufferSetImpl::set_matrix_inverse_beam( std::shared_ptr<std::vector<std::complex<float>>> matrix_inverse_beam) { m_matrix_inverse_beam = matrix_inverse_beam; - m_avg_aterm_correction = Array4D<std::complex<float>>( - m_matrix_inverse_beam->data(), m_subgridsize, m_subgridsize, 4, 4); - m_proxy->set_avg_aterm_correction(m_avg_aterm_correction); + auto matrix_inverse_beam_span = + aocommon::xt::CreateSpan<std::complex<float>, 4>( + m_matrix_inverse_beam->data(), {m_subgridsize, m_subgridsize, 4, 4}); + m_proxy->set_avg_aterm_correction(matrix_inverse_beam_span); } void BufferSetImpl::unset_matrix_inverse_beam() { m_matrix_inverse_beam.reset(); - m_avg_aterm_correction = Array4D<std::complex<float>>(0, 0, 0, 0); m_proxy->unset_avg_aterm_correction(); } diff --git a/idg-api/BufferSetImpl.h b/idg-api/BufferSetImpl.h index 44190a4d04e1144d33d99c0fb5caff76607c4fa1..9156eed25f1668ae64d7bc33cc11ad8b5c37b172 100644 --- a/idg-api/BufferSetImpl.h +++ b/idg-api/BufferSetImpl.h @@ -80,18 +80,14 @@ class BufferSetImpl : public virtual BufferSet { float get_w_step() const { return m_w_step; } const std::array<float, 2>& get_shift() const { return m_shift; } float get_kernel_size() const { return m_kernel_size; } - const Array2D<float>& get_spheroidal() const { return m_spheroidal; } + const aocommon::xt::Span<float, 2>& get_spheroidal() const { + return m_spheroidal; + } Stopwatch& get_watch(Watch watch) const; bool get_do_gridding() const { return m_do_gridding; } bool get_apply_aterm() const { return m_apply_aterm; } - const Array4D<std::complex<float>>& get_default_aterm_correction() const { - return m_default_aterm_correction; - } - const Array4D<std::complex<float>>& get_avg_aterm_correction() const { - return m_avg_aterm_correction; - } proxy::Proxy& get_proxy() const { return *m_proxy; } @@ -109,12 +105,10 @@ class BufferSetImpl : public virtual BufferSet { std::vector<float> m_taper_subgrid; std::vector<float> m_taper_grid; std::vector<float> m_inv_taper; - Array2D<float> m_spheroidal; + aocommon::xt::Span<float, 2> m_spheroidal; std::vector<std::complex<float>> m_average_beam; std::shared_ptr<std::vector<float>> m_scalar_beam; std::shared_ptr<std::vector<std::complex<float>>> m_matrix_inverse_beam; - Array4D<std::complex<float>> m_default_aterm_correction; - Array4D<std::complex<float>> m_avg_aterm_correction; bool m_stokes_I_only; int m_nr_correlations; int m_nr_polarizations; diff --git a/idg-api/BulkDegridder.cpp b/idg-api/BulkDegridder.cpp index f37d061c73332ab45d615761524557849481796f..cab3273d9055d5f83f05fbdd83d66195edafe000 100644 --- a/idg-api/BulkDegridder.cpp +++ b/idg-api/BulkDegridder.cpp @@ -20,11 +20,10 @@ BulkDegridderImpl::BulkDegridderImpl(const BufferSetImpl& bufferset, const std::size_t nr_stations) : bufferset_(bufferset), frequencies_( - bufferset.get_proxy().allocate_array1d<float>(frequencies.size())), + bufferset.get_proxy().allocate_span<float, 1>({frequencies.size()})), nr_stations_(nr_stations), - shift_(bufferset.get_shift().size()) { + shift_(bufferset.get_shift()) { std::copy_n(frequencies.data(), frequencies.size(), frequencies_.data()); - std::copy_n(bufferset.get_shift().data(), shift_.size(), shift_.data()); } BulkDegridderImpl::~BulkDegridderImpl() {} @@ -53,7 +52,7 @@ void BulkDegridderImpl::compute_visibilities( const size_t subgridsize = bufferset_.get_subgridsize(); proxy::Proxy& proxy = bufferset_.get_proxy(); - const int nr_correlations = bufferset_.get_nr_correlations(); + const size_t nr_correlations = bufferset_.get_nr_correlations(); static const double kDefaultUVWFactors[3] = {1.0, 1.0, 1.0}; if (!uvw_factors) uvw_factors = kDefaultUVWFactors; @@ -68,11 +67,13 @@ void BulkDegridderImpl::compute_visibilities( auto bufferUVW_tensor = proxy.allocate_tensor<UVW<float>, 2>({nr_baselines, nr_timesteps}); auto bufferUVW = bufferUVW_tensor.Span(); - auto bufferVisibilities = proxy.allocate_array4d<std::complex<float>>( - nr_baselines, nr_timesteps, frequencies_.size(), nr_correlations); + auto bufferVisibilities_tensor = + proxy.allocate_tensor<std::complex<float>, 4>( + {nr_baselines, nr_timesteps, frequencies_.size(), nr_correlations}); + auto bufferVisibilities = bufferVisibilities_tensor.Span(); // The proxy does not touch visibilities for out-of-bound uvw coordinates. // Since we copy all visibilities to the caller, initialize them to zero. - bufferVisibilities.zero(); + bufferVisibilities.fill(std::complex<float>(0, 0)); baseline_map.reserve(nr_baselines); for (size_t bl = 0; bl < nr_baselines; ++bl) { @@ -99,10 +100,8 @@ void BulkDegridderImpl::compute_visibilities( // The proxy expects the number of time steps as last value. std::vector<unsigned int> local_aterm_offsets = aterm_offsets; local_aterm_offsets.push_back(nr_timesteps); - const std::array<size_t, 1> local_aterm_offsets_shape{ - local_aterm_offsets.size()}; - auto aterm_offsets_span = aocommon::xt::CreateSpan(local_aterm_offsets.data(), - local_aterm_offsets_shape); + auto aterm_offsets_span = aocommon::xt::CreateSpan<unsigned int, 1>( + local_aterm_offsets, {local_aterm_offsets.size()}); // If aterms is empty, create default values and update the pointer. std::vector<Matrix2x2<std::complex<float>>> default_aterms; @@ -115,11 +114,10 @@ void BulkDegridderImpl::compute_visibilities( // The const cast is needed since proxy.degridding accepts const references // to arrays with non-const values, and the constructor for such arrays only // accepts non-const pointers. - auto* local_aterms = reinterpret_cast<Matrix2x2<std::complex<float>>*>( - const_cast<std::complex<float>*>(aterms)); - const Array4D<Matrix2x2<std::complex<float>>> aterms_array( - local_aterms, aterm_offsets_span.size() - 1, nr_stations_, subgridsize, - subgridsize); + using Aterm = Matrix2x2<std::complex<float>>; + auto aterms_span = aocommon::xt::CreateSpan<Aterm, 4>( + reinterpret_cast<Aterm*>(const_cast<std::complex<float>*>(aterms)), + {aterm_offsets_span.size() - 1, nr_stations_, subgridsize, subgridsize}); // Set Plan options Plan::Options options; @@ -131,18 +129,15 @@ void BulkDegridderImpl::compute_visibilities( // Create plan bufferset_.get_watch(BufferSetImpl::Watch::kPlan).Start(); - const std::array<size_t, 1> frequencies_shape{frequencies_.size()}; - auto frequencies_span = - aocommon::xt::CreateSpan(frequencies_.data(), frequencies_shape); std::unique_ptr<Plan> plan = - proxy.make_plan(bufferset_.get_kernel_size(), frequencies_span, bufferUVW, + proxy.make_plan(bufferset_.get_kernel_size(), frequencies_, bufferUVW, bufferStationPairs, aterm_offsets_span, options); bufferset_.get_watch(BufferSetImpl::Watch::kPlan).Pause(); // Run degridding bufferset_.get_watch(BufferSetImpl::Watch::kDegridding).Start(); proxy.degridding(*plan, frequencies_, bufferVisibilities, bufferUVW, - bufferStationPairs, aterms_array, aterm_offsets_span, + bufferStationPairs, aterms_span, aterm_offsets_span, bufferset_.get_spheroidal()); bufferset_.get_watch(BufferSetImpl::Watch::kDegridding).Pause(); @@ -156,8 +151,8 @@ void BulkDegridderImpl::compute_visibilities( for (size_t bl = 0; bl < nr_baselines; ++bl) { const int local_bl = baseline_map[bl]; if (local_bl != -1) { - const auto* in = reinterpret_cast<std::complex<float>*>( - bufferVisibilities.data(local_bl, t)); + const std::complex<float>* in = + &bufferVisibilities(local_bl, t, 0, 0); std::complex<float>* out = visibilities[t] + bl * baseline_size; std::copy_n(in, baseline_size, out); } @@ -172,8 +167,8 @@ void BulkDegridderImpl::compute_visibilities( for (size_t bl = 0; bl < nr_baselines; ++bl) { const int local_bl = baseline_map[bl]; if (local_bl != -1) { - const auto* in = reinterpret_cast<std::complex<float>*>( - bufferVisibilities.data(local_bl, t)); + const std::complex<float>* in = + &bufferVisibilities(local_bl, t, 0, 0); std::complex<float>* out = visibilities[t] + bl * baseline_size; for (size_t i = 0; i < nr_channels; ++i) { out[i * nr_correlations_out] = in[i * nr_correlations]; diff --git a/idg-api/BulkDegridderImpl.h b/idg-api/BulkDegridderImpl.h index 7a66af379fe7bd64e1ca00370c01346c64138fd6..a2359a4d86feabbcc6e5b92280268da8f02299bf 100644 --- a/idg-api/BulkDegridderImpl.h +++ b/idg-api/BulkDegridderImpl.h @@ -10,7 +10,7 @@ #include "BulkDegridder.h" -#include "common/ArrayTypes.h" +#include <aocommon/xt/span.h> namespace idg { namespace api { @@ -36,9 +36,9 @@ class BulkDegridderImpl : public BulkDegridder { private: const BufferSetImpl& bufferset_; - Array1D<float> frequencies_; + aocommon::xt::Span<float, 1> frequencies_; std::size_t nr_stations_; - Array1D<float> shift_; + std::array<float, 2> shift_; }; } // namespace api diff --git a/idg-api/DegridderBuffer.cpp b/idg-api/DegridderBuffer.cpp index 6c4605bfbdc081edafb364d64f9b2529d9f90bdd..e818e185aaf11ecbb297c63ea7bf4ac5224770cf 100644 --- a/idg-api/DegridderBuffer.cpp +++ b/idg-api/DegridderBuffer.cpp @@ -18,7 +18,9 @@ DegridderBufferImpl::DegridderBufferImpl(const BufferSetImpl& bufferset, size_t bufferTimesteps) : BufferImpl(bufferset, bufferTimesteps), m_buffer_full(false), - m_data_read(true) { + m_data_read(true), + m_bufferVisibilities2(aocommon::xt::CreateSpan<std::complex<float>, 4>( + nullptr, {0, 0, 0, 0})) { #if defined(DEBUG) cout << __func__ << endl; #endif @@ -75,8 +77,7 @@ bool DegridderBufferImpl::request_visibilities(size_t rowId, size_t timeIndex, // Keep mapping rowId -> (local_bl, local_time) for reading m_row_ids_to_data.emplace_back( - rowId, reinterpret_cast<std::complex<float>*>( - m_bufferVisibilities.data(local_bl, local_time))); + rowId, &m_bufferVisibilities(local_bl, local_time, 0, 0)); // Copy data into buffers m_bufferUVW(local_bl, local_time) = {static_cast<float>(uvwInMeters[0]), @@ -97,15 +98,12 @@ void DegridderBufferImpl::flush() { const size_t subgridsize = m_bufferset.get_subgridsize(); - m_aterm_offsets_array = - Array1D<unsigned int>(m_aterm_offsets.data(), m_aterm_offsets.size()); - m_aterms_array = Array4D<Matrix2x2<std::complex<float>>>( - m_aterms.data(), m_aterm_offsets_array.get_x_dim() - 1, m_nrStations, - subgridsize, subgridsize); - - const std::array<size_t, 1> aterm_offsets_shape{m_aterm_offsets_array.size()}; - auto aterm_offsets_span = aocommon::xt::CreateSpan( - m_aterm_offsets_array.data(), aterm_offsets_shape); + auto aterm_offsets_span = aocommon::xt::CreateSpan<unsigned int, 1>( + m_aterm_offsets.data(), {m_aterm_offsets.size()}); + auto aterms_span = + aocommon::xt::CreateSpan<Matrix2x2<std::complex<float>>, 4>( + m_aterms.data(), + {m_aterm_offsets.size() - 1, m_nrStations, subgridsize, subgridsize}); proxy::Proxy& proxy = m_bufferset.get_proxy(); @@ -119,24 +117,15 @@ void DegridderBufferImpl::flush() { // Create plan m_bufferset.get_watch(BufferSetImpl::Watch::kPlan).Start(); - const std::array<size_t, 1> frequencies_shape{m_frequencies.size()}; - auto frequencies_span = - aocommon::xt::CreateSpan(m_frequencies.data(), frequencies_shape); - const std::array<size_t, 1> station_pairs_shape{m_bufferStationPairs.size()}; - const std::array<size_t, 2> uvw_shape{m_bufferUVW.get_y_dim(), - m_bufferUVW.get_x_dim()}; - auto uvw_span = aocommon::xt::CreateSpan(m_bufferUVW.data(), uvw_shape); - auto station_pairs_span = aocommon::xt::CreateSpan( - m_bufferStationPairs.data(), station_pairs_shape); std::unique_ptr<Plan> plan = - proxy.make_plan(m_bufferset.get_kernel_size(), frequencies_span, uvw_span, - station_pairs_span, aterm_offsets_span, options); + proxy.make_plan(m_bufferset.get_kernel_size(), m_frequencies, m_bufferUVW, + m_bufferStationPairs, aterm_offsets_span, options); m_bufferset.get_watch(BufferSetImpl::Watch::kPlan).Pause(); // Run degridding m_bufferset.get_watch(BufferSetImpl::Watch::kDegridding).Start(); proxy.degridding(*plan, m_frequencies, m_bufferVisibilities, m_bufferUVW, - m_bufferStationPairs, m_aterms_array, m_aterm_offsets_array, + m_bufferStationPairs, aterms_span, aterm_offsets_span, m_bufferset.get_spheroidal()); m_bufferset.get_watch(BufferSetImpl::Watch::kDegridding).Pause(); @@ -175,15 +164,17 @@ DegridderBufferImpl::compute() { flush(); m_buffer_full = false; if (m_bufferset.get_nr_correlations() == 2) { - m_bufferVisibilities2.zero(); - for (int row_id = 0; row_id < m_row_ids_to_data.size(); ++row_id) { - for (int i = 0; i < m_nr_channels; ++i) { - m_bufferVisibilities2(row_id, i, 0) = - *(m_row_ids_to_data[row_id].second + i * 2); - m_bufferVisibilities2(row_id, i, 3) = - *(m_row_ids_to_data[row_id].second + i * 2 + 1); + m_bufferVisibilities2.fill(std::complex<float>(0.0f, 0.0f)); + for (size_t row_id = 0; row_id < m_row_ids_to_data.size(); ++row_id) { + const size_t bl = row_id / m_bufferTimesteps; + const size_t time = row_id % m_bufferTimesteps; + for (size_t chan = 0; chan < m_nr_channels; ++chan) { + m_bufferVisibilities2(bl, time, chan, 0) = + *(m_row_ids_to_data[row_id].second + chan * 2); + m_bufferVisibilities2(bl, time, chan, 3) = + *(m_row_ids_to_data[row_id].second + chan * 2 + 1); } - m_row_ids_to_data[row_id].second = m_bufferVisibilities2.data(row_id); + m_row_ids_to_data[row_id].second = &m_bufferVisibilities2(bl, time, 0, 0); } } return std::move(m_row_ids_to_data); @@ -202,8 +193,9 @@ void DegridderBufferImpl::malloc_buffers() { BufferImpl::malloc_buffers(); if (m_bufferset.get_nr_correlations() == 2) { - m_bufferVisibilities2 = Array3D<std::complex<float>>( - m_nr_baselines * m_bufferTimesteps, m_nr_channels, 4); + m_bufferVisibilities2 = + m_bufferset.get_proxy().allocate_span<std::complex<float>, 4>( + {m_nr_baselines, m_bufferTimesteps, m_nr_channels, 4}); } } diff --git a/idg-api/DegridderBufferImpl.h b/idg-api/DegridderBufferImpl.h index a0534c78fc1b7da9a7212f525512582665d8cbd5..deaa9493f9ae28a04462d4ec275a1e8efd097a8a 100644 --- a/idg-api/DegridderBufferImpl.h +++ b/idg-api/DegridderBufferImpl.h @@ -130,7 +130,8 @@ class DegridderBufferImpl : public virtual DegridderBuffer, public BufferImpl { bool m_buffer_full; bool m_data_read; std::vector<std::pair<size_t, std::complex<float>*>> m_row_ids_to_data; - Array3D<std::complex<float>> m_bufferVisibilities2; + aocommon::xt::Span<std::complex<float>, 4> + m_bufferVisibilities2; // BL x TI x CH x CR }; } // namespace api diff --git a/idg-api/GridderBuffer.cpp b/idg-api/GridderBuffer.cpp index b5f7a4cf3f237866c77fdcb3712fa02cdb1bd69f..f7d27c4837cd395023bac09983b8ed5117ed80f8 100644 --- a/idg-api/GridderBuffer.cpp +++ b/idg-api/GridderBuffer.cpp @@ -20,10 +20,17 @@ namespace api { GridderBufferImpl::GridderBufferImpl(const BufferSetImpl& bufferset, size_t bufferTimesteps) : BufferImpl(bufferset, bufferTimesteps), - m_bufferUVW2(0, 0), - m_bufferStationPairs2(0), - m_buffer_weights(0, 0, 0, 0), - m_buffer_weights2(0, 0, 0, 0), + m_bufferUVW2( + aocommon::xt::CreateSpan<idg::UVW<float>, 2>(nullptr, {0, 0})), + m_bufferStationPairs2( + aocommon::xt::CreateSpan<std::pair<unsigned int, unsigned int>, 1>( + nullptr, {0})), + m_bufferVisibilities2(aocommon::xt::CreateSpan<std::complex<float>, 4>( + nullptr, {0, 0, 0, 0})), + m_buffer_weights( + aocommon::xt::CreateSpan<float, 4>(nullptr, {0, 0, 0, 0})), + m_buffer_weights2( + aocommon::xt::CreateSpan<float, 4>(nullptr, {0, 0, 0, 0})), m_average_beam(nullptr) { #if defined(DEBUG) cout << __func__ << endl; @@ -99,24 +106,37 @@ void GridderBufferImpl::grid_visibilities( void GridderBufferImpl::compute_avg_beam() { m_bufferset.get_watch(BufferSetImpl::Watch::kAvgBeam).Start(); - const unsigned int subgrid_size = m_bufferset.get_subgridsize(); - const unsigned int nr_aterms = m_aterm_offsets2.size() - 1; - const unsigned int nr_antennas = m_nrStations; + const size_t subgrid_size = m_bufferset.get_subgridsize(); + const size_t nr_aterms = m_aterm_offsets2.size() - 1; + const size_t nr_antennas = m_nrStations; // average beam is always computed for all polarizations (for now) - const int nr_correlations = 4; - - Array4D<Matrix2x2<std::complex<float>>> aterms( - m_aterms2.data(), nr_aterms, nr_antennas, subgrid_size, subgrid_size); - Array1D<unsigned int> aterm_offsets(m_aterm_offsets2.data(), nr_aterms + 1); - idg::Array4D<std::complex<float>> average_beam(m_average_beam, subgrid_size, - subgrid_size, nr_correlations, - nr_correlations); + const size_t nr_correlations = 4; + + const std::array<size_t, 4> average_beam_shape{ + subgrid_size, subgrid_size, nr_correlations, nr_correlations}; + const std::array<size_t, 4> weights_shape{m_nr_baselines, m_bufferTimesteps, + m_nr_channels, nr_correlations}; + const std::array<size_t, 2> uvw_shape{m_nr_baselines, m_bufferTimesteps}; + const std::array<size_t, 1> station_pairs_shape{m_nr_baselines}; + const std::array<size_t, 4> aterms_shape{nr_aterms, nr_antennas, subgrid_size, + subgrid_size}; + const std::array<size_t, 1> aterm_offsets_shape{nr_aterms + 1}; + auto uvw = aocommon::xt::CreateSpan(m_bufferUVW2.data(), uvw_shape); + auto station_pairs = aocommon::xt::CreateSpan(m_bufferStationPairs2.data(), + station_pairs_shape); + auto aterms = aocommon::xt::CreateSpan(m_aterms2.data(), aterms_shape); + auto aterm_offsets = + aocommon::xt::CreateSpan(m_aterm_offsets2.data(), aterm_offsets_shape); + auto average_beam = + aocommon::xt::CreateSpan(m_average_beam, average_beam_shape); + auto weights = + aocommon::xt::CreateSpan(m_buffer_weights2.data(), weights_shape); proxy::Proxy& proxy = m_bufferset.get_proxy(); - proxy.compute_avg_beam(m_nrStations, get_frequencies_size(), m_bufferUVW2, - m_bufferStationPairs2, aterms, aterm_offsets, - m_buffer_weights2, average_beam); + proxy.compute_avg_beam(m_nrStations, get_frequencies_size(), uvw, + station_pairs, aterms, aterm_offsets, weights, + average_beam); m_bufferset.get_watch(BufferSetImpl::Watch::kAvgBeam).Pause(); @@ -131,17 +151,23 @@ void GridderBufferImpl::flush_thread_worker() { const size_t subgridsize = m_bufferset.get_subgridsize(); - if (!m_bufferset.get_apply_aterm()) { - m_aterm_offsets_array = Array1D<unsigned int>( - m_default_aterm_offsets.data(), m_default_aterm_offsets.size()); - m_aterms_array = Array4D<Matrix2x2<std::complex<float>>>( - m_default_aterms.data(), m_default_aterm_offsets.size() - 1, - m_nrStations, subgridsize, subgridsize); - } - - const std::array<size_t, 1> aterm_offsets_shape{m_aterm_offsets_array.size()}; - auto aterm_offsets_span = aocommon::xt::CreateSpan( - m_aterm_offsets_array.data(), aterm_offsets_shape); + auto aterm_offsets_span = + m_bufferset.get_apply_aterm() + ? aocommon::xt::CreateSpan<unsigned int, 1>(m_aterm_offsets2.data(), + {m_aterm_offsets2.size()}) + : aocommon::xt::CreateSpan<unsigned int, 1>( + m_default_aterm_offsets.data(), + {m_default_aterm_offsets.size()}); + + auto aterms_span = + m_bufferset.get_apply_aterm() + ? aocommon::xt::CreateSpan<Matrix2x2<std::complex<float>>, 4>( + m_aterms2.data(), {m_aterm_offsets2.size() - 1, m_nrStations, + subgridsize, subgridsize}) + : aocommon::xt::CreateSpan<Matrix2x2<std::complex<float>>, 4>( + m_default_aterms.data(), + {m_default_aterm_offsets.size() - 1, m_nrStations, subgridsize, + subgridsize}); proxy::Proxy& proxy = m_bufferset.get_proxy(); @@ -155,24 +181,15 @@ void GridderBufferImpl::flush_thread_worker() { // Create plan m_bufferset.get_watch(BufferSetImpl::Watch::kPlan).Start(); - const std::array<size_t, 1> frequencies_shape{m_frequencies.size()}; - auto frequencies_span = - aocommon::xt::CreateSpan(m_frequencies.data(), frequencies_shape); - const std::array<size_t, 1> station_pairs_shape{m_bufferStationPairs.size()}; - const std::array<size_t, 2> uvw_shape{m_bufferUVW2.get_y_dim(), - m_bufferUVW2.get_x_dim()}; - auto uvw_span = aocommon::xt::CreateSpan(m_bufferUVW2.data(), uvw_shape); - auto station_pairs_span = aocommon::xt::CreateSpan( - m_bufferStationPairs2.data(), station_pairs_shape); - std::unique_ptr<Plan> plan = - proxy.make_plan(m_bufferset.get_kernel_size(), frequencies_span, uvw_span, - station_pairs_span, aterm_offsets_span, options); + std::unique_ptr<Plan> plan = proxy.make_plan( + m_bufferset.get_kernel_size(), m_frequencies, m_bufferUVW2, + m_bufferStationPairs2, aterm_offsets_span, options); m_bufferset.get_watch(BufferSetImpl::Watch::kPlan).Pause(); // Run gridding m_bufferset.get_watch(BufferSetImpl::Watch::kGridding).Start(); proxy.gridding(*plan, m_frequencies, m_bufferVisibilities2, m_bufferUVW2, - m_bufferStationPairs2, m_aterms_array, m_aterm_offsets_array, + m_bufferStationPairs2, aterms_span, aterm_offsets_span, m_bufferset.get_spheroidal()); m_bufferset.get_watch(BufferSetImpl::Watch::kGridding).Pause(); } @@ -192,15 +209,8 @@ void GridderBufferImpl::flush() { std::swap(m_bufferVisibilities, m_bufferVisibilities2); std::swap(m_buffer_weights, m_buffer_weights2); std::swap(m_aterm_offsets, m_aterm_offsets2); - m_aterm_offsets_array = - Array1D<unsigned int>(m_aterm_offsets2.data(), m_aterm_offsets2.size()); std::swap(m_aterms, m_aterms2); - assert(m_aterms2.size() == (m_aterm_offsets_array.get_x_dim() - 1) * - m_nrStations * subgridsize * subgridsize); - m_aterms_array = Array4D<Matrix2x2<std::complex<float>>>( - m_aterms2.data(), m_aterm_offsets_array.get_x_dim() - 1, m_nrStations, - subgridsize, subgridsize); m_flush_thread = std::thread(&GridderBufferImpl::flush_thread_worker, this); @@ -248,22 +258,22 @@ void GridderBufferImpl::finished() { void GridderBufferImpl::malloc_buffers() { BufferImpl::malloc_buffers(); - int nr_correlations = m_bufferset.get_nr_correlations(); + const size_t nr_correlations = m_bufferset.get_nr_correlations(); proxy::Proxy& proxy = m_bufferset.get_proxy(); m_bufferUVW2 = - proxy.allocate_array2d<UVW<float>>(m_nr_baselines, m_bufferTimesteps); - m_bufferVisibilities2 = proxy.allocate_array4d<std::complex<float>>( - m_nr_baselines, m_bufferTimesteps, m_nr_channels, nr_correlations); + proxy.allocate_span<UVW<float>, 2>({m_nr_baselines, m_bufferTimesteps}); + m_bufferVisibilities2 = proxy.allocate_span<std::complex<float>, 4>( + {m_nr_baselines, m_bufferTimesteps, m_nr_channels, nr_correlations}); m_bufferStationPairs2 = - proxy.allocate_array1d<std::pair<unsigned int, unsigned int>>( - m_nr_baselines); - int nr_correlations_weights = 4; + proxy.allocate_span<std::pair<unsigned int, unsigned int>, 1>( + {m_nr_baselines}); + const size_t nr_correlations_weights = 4; m_buffer_weights = - proxy.allocate_array4d<float>(m_nr_baselines, m_bufferTimesteps, - m_nr_channels, nr_correlations_weights); + proxy.allocate_span<float, 4>({m_nr_baselines, m_bufferTimesteps, + m_nr_channels, nr_correlations_weights}); m_buffer_weights2 = - proxy.allocate_array4d<float>(m_nr_baselines, m_bufferTimesteps, - m_nr_channels, nr_correlations_weights); + proxy.allocate_span<float, 4>({m_nr_baselines, m_bufferTimesteps, + m_nr_channels, nr_correlations_weights}); } } // namespace api diff --git a/idg-api/GridderBufferImpl.h b/idg-api/GridderBufferImpl.h index 5bdacac2645ed45acb37e51aeb03a873576eb78e..263c0a5cf79f5714e485f7e2302431b91d773804 100644 --- a/idg-api/GridderBufferImpl.h +++ b/idg-api/GridderBufferImpl.h @@ -101,12 +101,16 @@ class GridderBufferImpl : public virtual GridderBuffer, public BufferImpl { private: // secondary buffers - Array2D<UVW<float>> m_bufferUVW2; // BL x TI - Array1D<std::pair<unsigned int, unsigned int>> m_bufferStationPairs2; // BL - Array4D<std::complex<float>> m_bufferVisibilities2; // BL x TI x CH x CR + aocommon::xt::Span<UVW<float>, 2> m_bufferUVW2; // BL x TI + aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1> + m_bufferStationPairs2; // BL + aocommon::xt::Span<std::complex<float>, 4> + m_bufferVisibilities2; // BL x TI x CH x CR std::vector<Matrix2x2<std::complex<float>>> m_aterms2; // ST x SB x SB - Array4D<float> m_buffer_weights; // BL x TI x NR_CHANNELS x NR_CORRELATIONS - Array4D<float> m_buffer_weights2; // BL x TI x NR_CHANNELS x NR_CORRELATIONS + aocommon::xt::Span<float, 4> + m_buffer_weights; // BL x TI x NR_CHANNELS x NR_CORRELATIONS + aocommon::xt::Span<float, 4> + m_buffer_weights2; // BL x TI x NR_CHANNELS x NR_CORRELATIONS std::vector<unsigned int> m_aterm_offsets2; std::thread m_flush_thread; diff --git a/idg-bin/tests/cxx/common/common.cpp b/idg-bin/tests/cxx/common/common.cpp index 28f9187c570e87412779e212507fc0791fffbed1..ecf6fe581f3b44d0520acd63627901568552eb73 100644 --- a/idg-bin/tests/cxx/common/common.cpp +++ b/idg-bin/tests/cxx/common/common.cpp @@ -253,18 +253,19 @@ int compare(idg::proxy::Proxy& proxy1, idg::proxy::Proxy& proxy2, float tol) { #endif #if TEST_AVERAGE_BEAM - idg::Array4D<std::complex<float>> average_beam(subgrid_size, subgrid_size, 4, - 4); - idg::Array4D<std::complex<float>> average_beam_ref(subgrid_size, subgrid_size, - 4, 4); - idg::Array4D<float> weights(nr_baselines, nr_timesteps, nr_channels, 4); - weights.init(1.0f); - average_beam.init(0.0f); - average_beam_ref.init(0.0f); + auto average_beam = proxy2.allocate_span<std::complex<float>, 4>( + {subgrid_size, subgrid_size, 4, 4}); + auto average_beam_ref = proxy1.allocate_span<std::complex<float>, 4>( + {subgrid_size, subgrid_size, 4, 4}); + auto weights = proxy1.allocate_span<float, 4>( + {nr_baselines, nr_timesteps, nr_channels, 4}); + weights.fill(1); + average_beam.fill(std::complex<float>(0, 0)); + average_beam_ref.fill(std::complex<float>(0, 0)); proxy1.compute_avg_beam(nr_stations, nr_channels, uvw, baselines, aterms, - aterm_offsets, weights, average_beam); - proxy2.compute_avg_beam(nr_stations, nr_channels, uvw, baselines, aterms, aterm_offsets, weights, average_beam_ref); + proxy2.compute_avg_beam(nr_stations, nr_channels, uvw, baselines, aterms, + aterm_offsets, weights, average_beam); float average_beam_error = get_accuracy(subgrid_size * subgrid_size * 4 * 4, (std::complex<float>*)average_beam.data(), diff --git a/idg-lib/src/CPU/common/CPU.cpp b/idg-lib/src/CPU/common/CPU.cpp index ef527a248f1cc9bfd32c3cf2908cfa017b87b260..0406cd516f29e76b956c526b1a56459147fbf629 100644 --- a/idg-lib/src/CPU/common/CPU.cpp +++ b/idg-lib/src/CPU/common/CPU.cpp @@ -687,26 +687,31 @@ void CPU::do_transform(DomainAtoDomainB direction) { void CPU::do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) { + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) { #if defined(DEBUG) std::cout << __func__ << std::endl; #endif - const unsigned int nr_aterms = aterm_offsets.size() - 1; - const unsigned int nr_baselines = baselines.get_x_dim(); - const unsigned int nr_timesteps = uvw.get_x_dim(); - const unsigned int subgrid_size = average_beam.get_w_dim(); - const unsigned int nr_polarizations = 4; + const size_t nr_aterms = aterm_offsets.size() - 1; + const size_t nr_baselines = baselines.size(); + assert(uvw.shape(0) == nr_baselines); + const size_t nr_timesteps = uvw.shape(1); + const size_t subgrid_size = average_beam.shape(0); + assert(average_beam.shape(1) == subgrid_size); + const size_t nr_polarizations = 4; m_report->initialize(); m_kernels->set_report(m_report); - auto* baselines_ptr = reinterpret_cast<idg::Baseline*>(baselines.data()); - auto* aterms_ptr = reinterpret_cast<std::complex<float>*>(aterms.data()); + auto* baselines_ptr = reinterpret_cast<const Baseline*>(baselines.data()); + auto* aterms_ptr = + reinterpret_cast<const std::complex<float>*>(aterms.data()); m_kernels->run_average_beam( nr_baselines, nr_antennas, nr_timesteps, nr_channels, nr_aterms, diff --git a/idg-lib/src/CPU/common/CPU.h b/idg-lib/src/CPU/common/CPU.h index 1207e51a2dec5ca9f1dd676aa0cc613c15ba91e9..e6627570d86657c03c78979b2ea40509dbc1c40e 100644 --- a/idg-lib/src/CPU/common/CPU.h +++ b/idg-lib/src/CPU/common/CPU.h @@ -97,11 +97,13 @@ class CPU : public Proxy { void do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) override; + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) override; protected: void init_wtiles(int grid_size, int subgrid_size, float image_size, diff --git a/idg-lib/src/CUDA/common/CUDA.h b/idg-lib/src/CUDA/common/CUDA.h index 18d9761583e2e2e655bb05fd4afff2cd5748e1dd..44487cd080a789f5a583ccb73b16197688f6de33 100644 --- a/idg-lib/src/CUDA/common/CUDA.h +++ b/idg-lib/src/CUDA/common/CUDA.h @@ -51,11 +51,13 @@ class CUDA : public Proxy { */ virtual void do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) override; + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) override; protected: void init_devices(); diff --git a/idg-lib/src/CUDA/common/routines/Beam.cpp b/idg-lib/src/CUDA/common/routines/Beam.cpp index 4f8a175cd6f07cdc19720898d7036ac7374c16ef..8d488cc19fd5bb33adab28e2f5977a971e9d5605 100644 --- a/idg-lib/src/CUDA/common/routines/Beam.cpp +++ b/idg-lib/src/CUDA/common/routines/Beam.cpp @@ -9,20 +9,24 @@ namespace cuda { void CUDA::do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) { + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) { #if defined(DEBUG) std::cout << "CUDA::" << __func__ << std::endl; #endif - const unsigned int nr_polarizations = 4; - const unsigned int nr_aterms = aterm_offsets.size() - 1; - const unsigned int nr_baselines = baselines.get_x_dim(); - const unsigned int nr_timesteps = uvw.get_x_dim(); - const unsigned int subgrid_size = average_beam.get_w_dim(); + const size_t nr_aterms = aterm_offsets.size() - 1; + const size_t nr_baselines = baselines.size(); + assert(uvw.shape(0) == nr_baselines); + const size_t nr_timesteps = uvw.shape(1); + const size_t subgrid_size = average_beam.shape(0); + assert(average_beam.shape(1) == subgrid_size); + const size_t nr_polarizations = 4; InstanceCUDA& device = get_device(0); cu::Context& context = device.get_context(); @@ -32,9 +36,13 @@ void CUDA::do_compute_avg_beam( device.set_report(m_report); // Allocate device memory - cu::DeviceMemory d_aterms(context, aterms.bytes()); - cu::DeviceMemory d_baselines(context, baselines.bytes()); - cu::DeviceMemory d_aterm_offsets(context, aterm_offsets.bytes()); + const size_t sizeof_aterms = aterms.size() * sizeof(*aterms.data()); + const size_t sizeof_baselines = baselines.size() * sizeof(*baselines.data()); + const size_t sizeof_aterm_offsets = + aterm_offsets.size() * sizeof(*aterm_offsets.data()); + cu::DeviceMemory d_aterms(context, sizeof_aterms); + cu::DeviceMemory d_baselines(context, sizeof_baselines); + cu::DeviceMemory d_aterm_offsets(context, sizeof_aterm_offsets); // The average beam is constructed in double-precision on the device. // After all baselines are processed (i.e. all contributions are added), // the data is copied to the host and there converted to single-precision. @@ -71,18 +79,18 @@ void CUDA::do_compute_avg_beam( // Initialize job data struct JobData { - unsigned int current_nr_baselines; - void* uvw_ptr; - void* weights_ptr; + size_t current_nr_baselines; + const idg::UVW<float>* uvw_ptr; + const float* weights_ptr; }; std::vector<JobData> jobs; - for (unsigned bl = 0; bl < nr_baselines; bl += jobsize) { + for (size_t bl = 0; bl < nr_baselines; bl += jobsize) { JobData job; - unsigned int first_bl = bl; - unsigned int last_bl = std::min(bl + jobsize, nr_baselines); + const size_t first_bl = bl; + const size_t last_bl = std::min(bl + jobsize, nr_baselines); job.current_nr_baselines = last_bl - first_bl; - job.uvw_ptr = uvw.data(first_bl, 0); - job.weights_ptr = weights.data(first_bl, 0, 0, 0); + job.uvw_ptr = &uvw(first_bl, 0); + job.weights_ptr = &weights(first_bl, 0, 0, 0); if (job.current_nr_baselines > 0) { jobs.push_back(job); } @@ -100,10 +108,10 @@ void CUDA::do_compute_avg_beam( cu::Stream& executestream = device.get_execute_stream(); // Copy static data - htodstream.memcpyHtoDAsync(d_aterms, aterms.data(), aterms.bytes()); - htodstream.memcpyHtoDAsync(d_baselines, baselines.data(), baselines.bytes()); + htodstream.memcpyHtoDAsync(d_aterms, aterms.data(), sizeof_aterms); + htodstream.memcpyHtoDAsync(d_baselines, baselines.data(), sizeof_baselines); htodstream.memcpyHtoDAsync(d_aterm_offsets, aterm_offsets.data(), - aterm_offsets.bytes()); + sizeof_aterm_offsets); // Initialize average beam d_average_beam.zero(htodstream); diff --git a/idg-lib/src/common/Proxy.cpp b/idg-lib/src/common/Proxy.cpp index 8b0306f40eb3abc6018992747bbbb9fe66e36b4a..db48cd283d2f37dca44bcce286fe2dc62c85fb89 100644 --- a/idg-lib/src/common/Proxy.cpp +++ b/idg-lib/src/common/Proxy.cpp @@ -9,7 +9,10 @@ namespace idg { namespace proxy { -Proxy::Proxy() { m_report.reset(new Report()); } +Proxy::Proxy() + : m_avg_aterm_correction(aocommon::xt::CreateSpan<std::complex<float>, 4>( + nullptr, {0, 0, 0, 0})), + m_report(std::make_shared<Report>()) {} Proxy::~Proxy() {} @@ -265,17 +268,14 @@ void Proxy::calibrate_update( void Proxy::calibrate_finish() { do_calibrate_finish(); } void Proxy::set_avg_aterm_correction( - const Array4D<std::complex<float>>& avg_aterm_correction) { - // check_dimensions_avg_aterm_correction(); - std::complex<float>* data = avg_aterm_correction.data(); - size_t size = - avg_aterm_correction.get_x_dim() * avg_aterm_correction.get_y_dim() * - avg_aterm_correction.get_z_dim() * avg_aterm_correction.get_w_dim(); - m_avg_aterm_correction.resize(size); - std::copy(data, data + size, m_avg_aterm_correction.begin()); + const aocommon::xt::Span<std::complex<float>, 4>& avg_aterm_correction) { + m_avg_aterm_correction = avg_aterm_correction; } -void Proxy::unset_avg_aterm_correction() { m_avg_aterm_correction.resize(0); } +void Proxy::unset_avg_aterm_correction() { + m_avg_aterm_correction = + aocommon::xt::CreateSpan<std::complex<float>, 4>(nullptr, {0, 0, 0, 0}); +} void Proxy::transform(DomainAtoDomainB direction) { do_transform(direction); } @@ -300,11 +300,13 @@ void Proxy::transform(DomainAtoDomainB direction, std::complex<float>* grid_ptr, void Proxy::compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) { + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) { #if defined(DEBUG) std::cout << __func__ << std::endl; #endif @@ -315,12 +317,14 @@ void Proxy::compute_avg_beam( void Proxy::do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam) { - average_beam.init(1.0f); + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam) { + average_beam.fill(std::complex<float>(1.0f, 0.0f)); } void Proxy::check_dimensions( diff --git a/idg-lib/src/common/Proxy.h b/idg-lib/src/common/Proxy.h index 48ae4779d522fb08a3c5e1ce45fe76872e83bb08..11a3bf99b14b41b20f3eb302f5a4fcfe1a83e042 100644 --- a/idg-lib/src/common/Proxy.h +++ b/idg-lib/src/common/Proxy.h @@ -249,11 +249,13 @@ class Proxy { */ virtual void compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam); + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam); //! Methods for querying and disabling Proxy capabilities bool supports_wstacking() { @@ -269,7 +271,7 @@ class Proxy { } void set_avg_aterm_correction( - const Array4D<std::complex<float>>& avg_aterm_correction); + const aocommon::xt::Span<std::complex<float>, 4>& avg_aterm_correction); void unset_avg_aterm_correction(); //! Methods for memory management @@ -468,11 +470,13 @@ class Proxy { virtual void do_compute_avg_beam( const unsigned int nr_antennas, const unsigned int nr_channels, - const Array2D<UVW<float>>& uvw_array, - const Array1D<std::pair<unsigned int, unsigned int>>& baselines, - const Array4D<Matrix2x2<std::complex<float>>>& aterms, - const Array1D<unsigned int>& aterm_offsets, const Array4D<float>& weights, - idg::Array4D<std::complex<float>>& average_beam); + const aocommon::xt::Span<UVW<float>, 2>& uvw, + const aocommon::xt::Span<std::pair<unsigned int, unsigned int>, 1>& + baselines, + const aocommon::xt::Span<Matrix2x2<std::complex<float>>, 4>& aterms, + const aocommon::xt::Span<unsigned int, 1>& aterm_offsets, + const aocommon::xt::Span<float, 4>& weights, + aocommon::xt::Span<std::complex<float>, 4>& average_beam); protected: void check_dimensions( @@ -505,7 +509,7 @@ class Proxy { const int nr_correlations = 4; - std::vector<std::complex<float>> m_avg_aterm_correction; + aocommon::xt::Span<std::complex<float>, 4> m_avg_aterm_correction; protected: virtual bool do_supports_wstacking() { return false; }