transformation_registry.h

Go to the documentation of this file.

#pragma once
 
#include <concepts>
#include <string>
#include <utility>
 
#include <sdfg/builder/structured_sdfg_builder.h>
#include <sdfg/transformations/recorder.h>
#include <sdfg/transformations/transformation.h>
#include <sdfg/transformations/transformation_schema.h>
 
#include <nlohmann/json.hpp>
 
#include <sdfg/einsum/einsum.h>
#include <sdfg/fp_math/transformations/fuse_fma.h>
#include <sdfg/fp_math/transformations/ozaki_contraction.h>
#include <sdfg/fp_math/transformations/quantize_contraction.h>
#include <sdfg/fp_math/transformations/unfuse_fma.h>
#include <sdfg/parallelization/transformations/loop_parallelization.h>
#include <sdfg/targets/gpu/gpu_mma_einsum_transform.h>
#include <sdfg/tiles/transformations/local_storage.h>
#include <sdfg/tiles/transformations/software_pipelining.h>
#include <sdfg/tiles/transformations/tile_fusion.h>
#include <sdfg/tiles/transformations/tile_vectorizer.h>
#include <sdfg/transformations/loop_distribute.h>
#include <sdfg/transformations/loop_interchange.h>
#include <sdfg/transformations/loop_peeling.h>
#include <sdfg/transformations/loop_shift.h>
#include <sdfg/transformations/loop_skewing.h>
#include <sdfg/transformations/loop_split.h>
#include <sdfg/transformations/loop_tiling.h>
#include <sdfg/transformations/map_collapse.h>
#include <sdfg/transformations/map_delinearize.h>
#include <sdfg/transformations/multi_level_tiling.h>
#include <sdfg/transformations/offloading/cuda_offload_transform.h>
#include <sdfg/transformations/offloading/cuda_parallelize_nested_map.h>
#include <sdfg/transformations/offloading/cuda_transform.h>
#include <sdfg/transformations/offloading/gpu_condition_propagation.h>
#include <sdfg/transformations/offloading/gpu_loop_reordering.h>
#include <sdfg/transformations/offloading/gpu_offload_nested_loop.h>
#include <sdfg/transformations/offloading/rocm_offload_transform.h>
#include <sdfg/transformations/offloading/rocm_parallelize_nested_map.h>
#include <sdfg/transformations/offloading/rocm_transform.h>
#include <sdfg/transformations/omp_transform.h>
#include <sdfg/transformations/redundant_zero_init_elimination.h>
#include <sdfg/transformations/stream_k.h>
#include <sdfg/transformations/trivial_loop_elimination.h>
#include <sdfg/transformations/unroll_transform.h>
#include <sdfg/transformations/vectorize_transform.h>
 
namespace sdfg {
namespace transformations {

template<typename T>
concept replayable_transformation_concept = transformation_concept<T> && requires(
                                                                             builder::StructuredSDFGBuilder& builder,
                                                                             const nlohmann::json& desc,
                                                                             const T& transformation,
                                                                             nlohmann::json& out
                                                                         ) {
    { T::from_json(builder, desc) } -> std::same_as<T>;
    { transformation.to_json(out) } -> std::same_as<void>;
};
 
namespace detail {

template<typename T, typename Visitor>
    requires replayable_transformation_concept<T>
decltype(auto) invoke_for(Visitor&& visitor) {
    return std::forward<Visitor>(visitor).template operator()<T>();
}
 
} // namespace detail

template<typename Visitor>
decltype(auto)
dispatch_transformation(const std::string& transformation_name, const nlohmann::json& desc, Visitor&& visitor) {
    if (transformation_name == "LoopTiling") {
        return detail::invoke_for<transformations::LoopTiling>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "MapCollapse") {
        return detail::invoke_for<transformations::MapCollapse>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "MapDelinearize") {
        return detail::invoke_for<transformations::MapDelinearize>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "MultiLevelTiling") {
        return detail::invoke_for<transformations::MultiLevelTiling>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopDistribute") {
        return detail::invoke_for<transformations::LoopDistribute>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopInterchange") {
        return detail::invoke_for<transformations::LoopInterchange>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LocalStorage") {
        return detail::invoke_for<transformations::LocalStorage>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "SoftwarePipelining") {
        return detail::invoke_for<transformations::SoftwarePipelining>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "TileFusion") {
        return detail::invoke_for<transformations::TileFusion>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "TileVectorizer") {
        return detail::invoke_for<transformations::TileVectorizer>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopSkewing") {
        return detail::invoke_for<transformations::LoopSkewing>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopShift") {
        return detail::invoke_for<transformations::LoopShift>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopSplit") {
        return detail::invoke_for<transformations::LoopSplit>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopParallelization") {
        return detail::invoke_for<transformations::LoopParallelization>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "OMPTransform") {
        return detail::invoke_for<transformations::OMPTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "LoopPeeling") {
        return detail::invoke_for<transformations::LoopPeeling>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "TrivialLoopElimination") {
        return detail::invoke_for<transformations::TrivialLoopElimination>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "RedundantZeroInitElimination") {
        return detail::invoke_for<transformations::RedundantZeroInitElimination>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "StreamK") {
        return detail::invoke_for<transformations::StreamK>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "VectorizeTransform") {
        return detail::invoke_for<transformations::VectorizeTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "UnrollTransform") {
        return detail::invoke_for<transformations::UnrollTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "CUDATransform") {
        return detail::invoke_for<cuda::CUDATransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "CUDAOffloadTransform") {
        return detail::invoke_for<cuda::CUDAOffloadTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "ROCMOffloadTransform") {
        return detail::invoke_for<rocm::ROCMOffloadTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "CUDAParallelizeNestedMap") {
        return detail::invoke_for<transformations::CUDAParallelizeNestedMap>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "ROCMTransform") {
        return detail::invoke_for<rocm::ROCMTransform>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "ROCMParallelizeNestedMap") {
        return detail::invoke_for<transformations::ROCMParallelizeNestedMap>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "GPUOffloadNestedLoop") {
        const std::string gpu_type = desc.contains("parameters") && desc["parameters"].contains("gpu_type")
                                         ? desc["parameters"]["gpu_type"].get<std::string>()
                                         : cuda::ScheduleType_CUDA_Offload::value();
        if (gpu_type == rocm::ScheduleType_ROCM_Offload::value()) {
            return detail::invoke_for<
                transformations::GPUOffloadNestedLoop<rocm::ScheduleType_ROCM_Offload>>(std::forward<Visitor>(visitor));
        } else {
            return detail::invoke_for<
                transformations::GPUOffloadNestedLoop<cuda::ScheduleType_CUDA_Offload>>(std::forward<Visitor>(visitor));
        }
    } else if (transformation_name == "GPUConditionPropagation") {
        return detail::invoke_for<transformations::GPUConditionPropagation>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "GPUConditionPropagationScope") {
        return detail::invoke_for<transformations::GPUConditionPropagationScope>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "GPULoopReordering") {
        return detail::invoke_for<transformations::GPULoopReordering>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "EinsumLift") {
        return detail::invoke_for<einsum::EinsumLift>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "EinsumExtend") {
        return detail::invoke_for<einsum::EinsumExtend>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "EinsumExpand") {
        return detail::invoke_for<einsum::EinsumPromotion>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "Einsum2Dot") {
        return detail::invoke_for<einsum::Einsum2Dot>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "Einsum2Gemm") {
        return detail::invoke_for<einsum::Einsum2Gemm>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "FuseFMA") {
        return detail::invoke_for<fp_math::FuseFMA>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "UnfuseFMA") {
        return detail::invoke_for<fp_math::UnfuseFMA>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "QuantizeContraction") {
        return detail::invoke_for<fp_math::QuantizeContraction>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "OzakiContraction") {
        return detail::invoke_for<fp_math::OzakiContraction>(std::forward<Visitor>(visitor));
    } else if (transformation_name == "GpuMmaEinsumTransform") {
        return detail::invoke_for<gpu::GpuMmaEinsumTransform>(std::forward<Visitor>(visitor));
    } else {
        throw transformations::InvalidTransformationDescriptionException("Unknown transformation: " + transformation_name);
    }
}

template<typename T>
    requires replayable_transformation_concept<T>
bool verify_round_trip(builder::StructuredSDFGBuilder& builder, const nlohmann::json& desc, std::string& error_out) {
    // 1) Reconstruct from the (possibly enriched) description and re-serialize.
    T first(T::from_json(builder, desc));
    nlohmann::json canonical;
    first.to_json(canonical);
 
    // 2) The output of to_json must itself satisfy the shared schema.
    if (!validate_transformation_schema(canonical, error_out)) {
        error_out = "to_json() produced a schema-invalid description: " + error_out;
        return false;
    }
 
    // 3) The reconstructed type must match the description that selected it.
    if (desc.contains("transformation_type") && desc.at("transformation_type") != canonical.at("transformation_type")) {
        error_out = "transformation_type changed across round-trip: description '" +
                    desc.at("transformation_type").get<std::string>() + "' vs to_json '" +
                    canonical.at("transformation_type").get<std::string>() + "'";
        return false;
    }
 
    // 4) from_json/to_json must be a stable fixed point: any field written by
    //    to_json but not read back by from_json would perturb this comparison.
    T second(T::from_json(builder, canonical));
    nlohmann::json canonical_again;
    second.to_json(canonical_again);
    if (canonical != canonical_again) {
        error_out = "from_json/to_json is not a stable round-trip: " + canonical.dump() +
                    " != " + canonical_again.dump();
        return false;
    }
 
    return true;
}

inline bool verify_transformation_round_trip(
    builder::StructuredSDFGBuilder& builder, const nlohmann::json& desc, std::string& error_out
) {
    if (!desc.contains("transformation_type") || !desc.at("transformation_type").is_string()) {
        error_out = "description is missing a string 'transformation_type'";
        return false;
    }
    const auto transformation_name = desc.at("transformation_type").get<std::string>();
    std::string local_error;
    const bool ok = dispatch_transformation(transformation_name, desc, [&]<typename T>() -> bool {
        return verify_round_trip<T>(builder, desc, local_error);
    });
    if (!ok) {
        error_out = std::move(local_error);
    }
    return ok;
}
 
} // namespace transformations
} // namespace sdfg