op_dense_flow()

bool MayaFlux::Yantra::VisionGpuExecutor::op_dense_flow ( VisionGpuContexts &  contexts, const Kinesis::Vision::VisionStep &  step  )

private

Dense optical flow from the previous frame to the current one.

Needs an earlier RgbaToGray step, whose gray frame the pyramid hook has already turned into the current atlas. The whole solve is one dependency sequence: for each level from coarsest to finest, the previous frame's structure tensor and then a number of warp, box sum and solve iterations. The flow of a finer level starts from the coarser one upsampled and doubled, seeded inside the first warp of the level.

The first frame after a reset has no previous frame, so it only ends the frame and produces no flow.

Returns

True when the step was deferred and the run must suspend.

Definition at line 503 of file VisionFlowDispatch.cpp.

{
    auto& state = contexts.flow_state;
    auto& flow = contexts.flow;
    auto& foundry = Portal::Graphics::get_shader_foundry();
 
    if (!state.curr_ready || !state.atlas[0]) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
            "run_gpu: OpticalFlowDense requires an earlier RgbaToGray step to supply the gray frame");
        return false;
    }
 
    const auto& p = std::get<OpticalFlowDenseParams>(step.params);
    const auto& layout = state.layout;
    const uint32_t w = contexts.pass.w;
    const uint32_t h = contexts.pass.h;
 
    if (w != layout.level[0].w || h != layout.level[0].h) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
            "run_gpu: OpticalFlowDense image is {}x{} but the gray frame was captured at {}x{}",
            w, h, layout.level[0].w, layout.level[0].h);
        return false;
    }
 
    if (!state.have_prev) {
        commit_flow_frame(contexts);
        return false;
    }
 
    if (!ensure_dense_images(contexts)) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
            "run_gpu: OpticalFlowDense failed to allocate its flow images");
        return false;
    }
 
    struct DenseImages {
        std::shared_ptr<Core::VKImage> atlas_a;
        std::shared_ptr<Core::VKImage> atlas_b;
        std::shared_ptr<Core::VKImage> flow_out_a;
        std::shared_ptr<Core::VKImage> flow_out_b;
        std::shared_ptr<Core::VKImage> flow_lvl;
        std::shared_ptr<Core::VKImage> scratch_a;
        std::shared_ptr<Core::VKImage> scratch_b;
        std::shared_ptr<Core::VKImage> tensor;
        std::shared_ptr<Core::VKImage> vis;
    };
 
    const auto images = std::make_shared<const DenseImages>(DenseImages {
        .atlas_a = state.atlas[0],
        .atlas_b = state.atlas[1],
        .flow_out_a = state.flow_out[0],
        .flow_out_b = state.flow_out[1],
        .flow_lvl = state.flow_lvl,
        .scratch_a = state.dense_a,
        .scratch_b = state.dense_b,
        .tensor = state.dense_tensor,
        .vis = state.flow_vis,
    });
 
    const uint32_t out_parity = state.curr;
    const uint32_t radius = std::clamp(p.window_radius, 1U, k_flow_dense_max_radius);
    const uint32_t iterations = std::max(p.iterations, 1U);
 
    const auto image_hazard = [](const std::shared_ptr<Core::VKImage>& image, uint32_t binding) {
        return HazardResource {
            .binding = { .set = 0, .binding = binding, .direction = GpuBufferBinding::Direction::INPUT_OUTPUT, .element_type = GpuBufferBinding::ElementType::IMAGE_STORAGE },
            .image = image->get_image(),
        };
    };
 
    const auto scratch_a_hazard = image_hazard(images->scratch_a, 8U);
    const auto scratch_b_hazard = image_hazard(images->scratch_b, 9U);
    const auto tensor_hazard = image_hazard(images->tensor, 10U);
    const auto coarse_hazard = image_hazard(images->flow_lvl, 7U);
    const auto out_hazard = image_hazard(out_parity == 0 ? images->flow_out_a : images->flow_out_b, 5U + out_parity);
 
    std::vector<DependencyStage> stages;
    stages.reserve(static_cast<size_t>(layout.levels) * (2U + 3U * iterations) + 1U);
 
    const auto push_stage = [&](const FlowDensePC& stage_pc, const PyramidLevel& region, std::vector<HazardResource> hazards) {
        stages.push_back({
            .config = { .shader_path = "flow_dense.comp.spv", .workgroup_size = k_wg2d, .push_constant_size = sizeof(FlowDensePC) },
            .stage_fn = [images, stage_pc](GpuDispatchCore& ctx) {
                    using Kind = GpuBufferBinding::ElementType;
                    ctx.stage_image_at(3, images->atlas_a, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(4, images->atlas_b, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(5, images->flow_out_a, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(6, images->flow_out_b, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(7, images->flow_lvl, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(8, images->scratch_a, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(9, images->scratch_b, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(10, images->tensor, Kind::IMAGE_STORAGE);
                    ctx.stage_image_at(11, images->vis, Kind::IMAGE_STORAGE);
                    ctx.set_push_constants(stage_pc); },
            .hazard_fn = [hazards = std::move(hazards)](GpuDispatchCore&) { return hazards; },
            .explicit_groups = std::array<uint32_t, 3> { (region.w + k_wg2d[0] - 1U) / k_wg2d[0], (region.h + k_wg2d[1] - 1U) / k_wg2d[1], 1U },
        });
    };
 
    for (uint32_t level = layout.levels; level-- > 0;) {
        const PyramidLevel lv = layout.level[level];
        const bool coarsest = level + 1U == layout.levels;
        const PyramidLevel coarse = coarsest ? PyramidLevel {} : layout.level[level + 1];
        const auto& flow_hazard = level == 0 ? out_hazard : coarse_hazard;
 
        FlowDensePC pc {
            .phase = 0,
            .curr_atlas = state.curr,
            .out_parity = out_parity,
            .level = level,
            .lvl_ox = lv.ox,
            .lvl_oy = lv.oy,
            .lvl_w = lv.w,
            .lvl_h = lv.h,
            .crs_ox = coarse.ox,
            .crs_oy = coarse.oy,
            .crs_w = coarse.w,
            .crs_h = coarse.h,
            .has_coarse = coarsest ? 0U : 1U,
            .radius = radius,
            .eigen_threshold = p.eigen_threshold,
            .max_step = p.max_step,
            .visual_range = p.visual_range,
            .visual_min_motion = p.visual_min_motion,
            .last_iteration = 0,
        };
 
        const auto add = [&](FlowDensePhase phase, std::vector<HazardResource> hazards) {
            pc.phase = static_cast<uint32_t>(phase);
            push_stage(pc, lv, std::move(hazards));
        };
 
        add(FlowDensePhase::TENSOR_H, { scratch_a_hazard });
        add(FlowDensePhase::TENSOR_V, { tensor_hazard, scratch_a_hazard });
        for (uint32_t iteration = 0; iteration < iterations; ++iteration) {
            add(iteration == 0 ? FlowDensePhase::WARP_INIT : FlowDensePhase::WARP, { scratch_a_hazard, flow_hazard });
            add(FlowDensePhase::BOX_H, { scratch_b_hazard, scratch_a_hazard });
            pc.last_iteration = iteration + 1U == iterations ? 1U : 0U;
            add(FlowDensePhase::SOLVE, { flow_hazard, scratch_b_hazard });
        }
    }
 
    if (p.visualize) {
        const PyramidLevel finest = layout.level[0];
        const FlowDensePC vis_pc {
            .phase = static_cast<uint32_t>(FlowDensePhase::VISUALIZE),
            .curr_atlas = state.curr,
            .out_parity = out_parity,
            .level = 0,
            .lvl_ox = finest.ox,
            .lvl_oy = finest.oy,
            .lvl_w = finest.w,
            .lvl_h = finest.h,
            .crs_ox = 0,
            .crs_oy = 0,
            .crs_w = 0,
            .crs_h = 0,
            .has_coarse = 0,
            .radius = radius,
            .eigen_threshold = p.eigen_threshold,
            .max_step = p.max_step,
            .visual_range = p.visual_range,
            .visual_min_motion = p.visual_min_motion,
            .last_iteration = 0,
        };
        push_stage(vis_pc, finest, { out_hazard, image_hazard(images->vis, 11U) });
    }
 
    ExecutionContext dense_ctx;
    dense_ctx.mode = ExecutionMode::DEPENDENCY;
    DependencyParams dense_params;
    dense_params.stages = stages;
    dense_params.async = true;
    dense_ctx.parameters = dense_params;
    const auto dense_dispatch_result = flow.execute(Datum<> {}, dense_ctx);
    const auto fence = dense_dispatch_result.get_metadata<Portal::Graphics::FenceID>("gpu_fence").value_or(Portal::Graphics::INVALID_FENCE);
    if (fence == Portal::Graphics::INVALID_FENCE) {
        MF_ERROR(Journal::Component::Yantra, Journal::Context::ComputeMatrix,
            "run_gpu: OpticalFlowDense failed to submit its dispatch sequence");
        return false;
    }
 
    if (step.deferred) {
        contexts.suspended.fence = fence;
        contexts.suspended.finalize = [visualize = p.visualize](VisionGpuContexts& c) { finish_dense(c, visualize); };
        return true;
    }
 
    foundry.wait_for_fence(fence);
    foundry.release_fence(fence);
    finish_dense(contexts, p.visualize);
    return false;
}

References MayaFlux::Yantra::DependencyParams::async, MayaFlux::Portal::Graphics::HazardResource::binding, MayaFlux::Journal::ComputeMatrix, MayaFlux::Kinesis::Vision::VisionStep::deferred, MayaFlux::Yantra::DEPENDENCY, fence, MayaFlux::Yantra::VisionGpuContexts::Suspension::fence, MayaFlux::Yantra::VisionGpuContexts::Suspension::finalize, MayaFlux::Yantra::VisionGpuContexts::flow, MayaFlux::Yantra::VisionGpuContexts::flow_state, MayaFlux::Portal::Graphics::get_shader_foundry(), MayaFlux::Kinesis::Vision::PyramidLevel::h, MayaFlux::Kinesis::Vision::VisionPass< Handle >::h, h, image, MayaFlux::Portal::Graphics::GpuBufferBinding::IMAGE_STORAGE, MayaFlux::Portal::Graphics::GpuBufferBinding::INPUT_OUTPUT, MayaFlux::Portal::Graphics::INVALID_FENCE, MayaFlux::Yantra::VisionInternal::k_flow_dense_max_radius, MayaFlux::Yantra::VisionInternal::k_wg2d, MF_ERROR, MayaFlux::Yantra::ExecutionContext::mode, MayaFlux::Kinesis::Vision::PyramidLevel::ox, MayaFlux::Kinesis::Vision::PyramidLevel::oy, MayaFlux::Yantra::ExecutionContext::parameters, MayaFlux::Kinesis::Vision::VisionStep::params, MayaFlux::Yantra::VisionGpuContexts::pass, MayaFlux::Yantra::VisionInternal::FlowDensePC::phase, radius, MayaFlux::Portal::Graphics::GpuBufferBinding::set, MayaFlux::Yantra::DependencyParams::stages, MayaFlux::Yantra::VisionGpuContexts::suspended, MayaFlux::Kinesis::Vision::PyramidLevel::w, MayaFlux::Kinesis::Vision::VisionPass< Handle >::w, and MayaFlux::Journal::Yantra.

Here is the call graph for this function: