sort_implementation()

Datum< Kakshya::RegionGroup > MayaFlux::Yantra::Granular::SortOp::sort_implementation ( const Datum< Kakshya::RegionGroup > &  input )

overrideprotected

Definition at line 617 of file GranularWorkflow.cpp.

{
    const auto feature_key = this->template get_parameter_or_default<std::string>("feature_key", std::string("feature"));
    const bool ascending = this->template get_parameter_or_default<bool>("ascending", true);
    const auto gpu_thresh = this->template get_parameter_or_default<uint32_t>("gpu_sort_threshold", 0U);
    const auto n = static_cast<uint32_t>(input.data.regions.size());
 
    Datum<Kakshya::RegionGroup> out = input;
 
    if (n == 0)
        return out;
 
    if (gpu_thresh > 0 && n >= gpu_thresh) {
        struct SortPC {
            uint32_t element_count;
            uint32_t pass_number;
            uint32_t ascending;
        };
 
        std::vector<float> values(n);
        for (uint32_t i = 0; i < n; ++i) {
            auto attr = out.data.regions[i].get_attribute<double>(feature_key);
            values[i] = attr ? static_cast<float>(*attr) : 0.0F;
        }
 
        std::vector<uint32_t> indices(n);
        std::iota(indices.begin(), indices.end(), 0U);
 
        auto executor = std::make_shared<ShaderExecutionContext<>>(
            GpuShaderConfig {
                .shader_path = "sort_by_attribute.comp",
                .workgroup_size = { 256, 1, 1 },
                .push_constant_size = sizeof(SortPC) });
 
        executor->in_out(0, values, GpuBufferBinding::ElementType::FLOAT32)
            .in_out(1, indices, GpuBufferBinding::ElementType::UINT32)
            .set_output_size(1, n * sizeof(uint32_t));
 
        const uint32_t n_passes = n * 2;
        executor->set_multipass(n_passes,
            [n, ascending](uint32_t pass, void* pc_data) {
                const SortPC pc {
                    .element_count = n,
                    .pass_number = pass & 1U,
                    .ascending = ascending ? 1U : 0U,
                };
                std::memcpy(pc_data, &pc, sizeof(SortPC));
            });
 
        auto gpu_sorter = std::make_shared<GpuSorter<>>(executor);
 
        Datum<std::vector<Kakshya::DataVariant>> sort_input {
            { Kakshya::DataVariant(std::vector<double>(values.begin(), values.end())) }
        };
 
        auto result = gpu_sorter->apply_operation(sort_input);
        const auto sorted_indices = ShaderExecutionContext<>::read_output<uint32_t>(result, 1);
 
        std::vector<Kakshya::Region> reordered(n);
        for (uint32_t i = 0; i < n; ++i)
            reordered[i] = out.data.regions[sorted_indices[i]];
 
        out.data.regions = std::move(reordered);
        out.data.current_region_index = 0;
        out.data.active_indices.clear();
        return out;
    }
 
    out.data.sort_by_attribute(feature_key);
 
    if (!ascending)
        std::ranges::reverse(out.data.regions);
 
    out.data.current_region_index = 0;
    out.data.active_indices.clear();
 
    return out;
}

References MayaFlux::Yantra::Datum< T >::data, MayaFlux::Yantra::GpuBufferBinding::FLOAT32, MayaFlux::Yantra::GpuShaderConfig::shader_path, and MayaFlux::Yantra::GpuBufferBinding::UINT32.