rocm_arch.h

Go to the documentation of this file.

#pragma once
 
#include <string>
#include <vector>
 
#include "sdfg/codegen/dispatchers/block_dispatcher.h"
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/targets/gpu/gpu_arch.h"
#include "sdfg/types/array.h"
 
namespace sdfg::gpu::rocm {
 
enum class RocmMmaVariant {
    Fragment = 1, 
    Bare = 2, 
};
 
struct RocmMmaSupport : public GpuMmaSupport {
    const bool base_support;
    const bool f32_support;
    const uint16_t threads_per_mma_block;
    const uint16_t int8_block_k;
 
    RocmMmaSupport(bool base_support, bool f32_support, uint16_t threads, uint16_t int8_block_k = 0)
        : GpuMmaSupport(), base_support(base_support), f32_support(f32_support), threads_per_mma_block(threads),
          int8_block_k(int8_block_k) {
    }

    MmaBlockSize int8_block_size() const {
        return {16, 16, int8_block_k};
    }
 
protected:
    static constexpr MmaBlockSize DEFAULT_BLOCK_SIZE = {16, 16, 16};
    static constexpr MmaBlockSize CDNA_BLOCK_SIZE = {32, 32, 8};
 
public:
    static constexpr const char* MMA_FRAGMENT_STORAGE_TYPE = "ROCM_MMA";
    static constexpr const char* MMA_REGS_STORAGE_TYPE = "ROCM_MMA_REGS";
 
    bool supported_types(types::PrimitiveType input_type, types::PrimitiveType output_type) const override;
 
    types::PrimitiveType get_accumulator_type(
        types::PrimitiveType input_type, types::PrimitiveType output_type, types::PrimitiveType desired_acc_type
    ) const override;
 
    bool is_valid_block_size(
        const MmaBlockSize& block_size, types::PrimitiveType input_type, types::PrimitiveType acc_type
    ) const override;
 
    std::optional<GpuMmaTiling> try_get_mma_tiling(
        const MmaBlockSize& block_size,
        const BlockCounts* wave_tile_hint,
        const symbolic::MultiExpression& res_shape,
        types::PrimitiveType acc_type,
        bool ceil_counts = false
    ) const;
 
    std::optional<GpuMmaTiling> get_mma_tiling(
        const symbolic::MultiExpression& res_shape,
        types::PrimitiveType input_type,
        types::PrimitiveType acc_type,
        const MmaBlockSize* block_size_hint,
        const BlockCounts* wave_tile_hint
    ) const override;
 
    std::optional<GpuMmaTiling> get_ragged_mma_tiling(
        const symbolic::MultiExpression& padded_shape,
        types::PrimitiveType input_type,
        types::PrimitiveType output_type,
        types::PrimitiveType acc_type,
        const MmaBlockSize* block_size_hint,
        const BlockCounts* wave_tile_hint
    ) const override;
 
    void set_mma_fragment_storage_type(
        types::StorageType& storage_type, const GpuMmaTiling& size, MmaFragmentType type, MmaFragmentLayout layout
    ) const override;
 
    static bool is_mma_type(const types::StorageType& storage);
 
    static RocmMmaVariant select_variant(const GpuMmaTiling& tiling);
 
    data_flow::ImplementationType get_mma_impl_type(const GpuMmaTiling& tiling) const override;
 
    void emit_block_frag_type(
        std::ostream& os, const types::StorageType& storage_type, types::PrimitiveType element_type
    ) const;
 
    static void emit_bare_fragment_decl_type(
        std::ostream& os,
        MmaFragmentType type,
        const MmaBlockSize& dims,
        MmaFragmentLayout layout,
        types::PrimitiveType scalar_type,
        std::optional<std::pair<int, int>> coop_dims
    );
 
    static void emit_block_frag_type(
        std::ostream& os,
        MmaFragmentType type,
        const MmaBlockSize& dims,
        MmaFragmentLayout layout,
        types::PrimitiveType scalar_type,
        std::optional<std::pair<int, int>> coop_dims
    );
 
    std::vector<MmaBlockSize>
    get_supported_block_sizes(types::PrimitiveType input_type, types::PrimitiveType acc_type) const override;
};
 
 
class RocmArch : public GpuArch {
    int per_cu_threads_;
    std::string rocm_name_;
    RocmMmaSupport mma_support_;
 
public:
    RocmArch(
        const std::string& name,
        int per_cu_threads,
        bool mma_base_support,
        bool mma_f32_support,
        uint16_t mma_int8_block_k = 0
    )
        : GpuArch(), rocm_name_(name), per_cu_threads_(per_cu_threads),
          mma_support_(mma_base_support ? 16 : 0, mma_f32_support, per_cu_threads, mma_int8_block_k) {
    }
 
    std::string unique_id() const override;
    std::string name() const override;
 
    int per_cu_threads() const override {
        return per_cu_threads_;
    }
 
    const RocmMmaSupport* mma_support() const override {
        if (mma_support_.base_support) {
            return &mma_support_;
        } else {
            return nullptr;
        }
    }
 
    structured_control_flow::ScheduleType create_schedule_type() const override;
};
 
extern RocmArch ROCM_ARCH_GFX1201;
 
extern RocmArch ROCM_ARCH_GFX90A;
extern RocmArch ROCM_ARCH_GFX942;

struct RocmDeviceInfo {
    std::string gfx_name;
    std::vector<std::string> device_names;
    int wavefront_size = 0;
};

const std::vector<RocmDeviceInfo>& query_rocm_devices();

const RocmArch* rocm_arch_for_device(const RocmDeviceInfo& device);

const RocmArch* rocm_arch_from_schedule_type(const structured_control_flow::ScheduleType& schedule);
const RocmArch* rocm_arch_from_env();
const RocmArch* rocm_arch_from_available_hardware();
const RocmArch* rocm_arch_parse(const std::string& name);
 
} // namespace sdfg::gpu::rocm