cuda_arch.h

Go to the documentation of this file.

#pragma once
 
#include "sdfg/structured_control_flow/structured_loop.h"
#include "sdfg/targets/gpu/gpu_arch.h"
#include "sdfg/types/array.h"
 
namespace sdfg::codegen {
class LanguageExtension;
}
 
namespace sdfg::gpu::cuda {

struct CudaComputeCapability {
    uint32_t compute_cap;
    std::vector<std::string> device_names;
};

std::vector<CudaComputeCapability> query_cuda_compute_capabilities();
 
struct CudaMmaSupport : public GpuMmaSupport {
    const bool base_support;
    const bool tf32_support;
    const bool fp64_support;
 
    CudaMmaSupport(bool base_support, bool tf32_support, bool fp64_support)
        : GpuMmaSupport(), base_support(base_support), tf32_support(tf32_support), fp64_support(fp64_support) {
    }
 
protected:
    static constexpr MmaBlockSize DEFAULT_BLOCK_SIZE = {16, 16, 16};
    static constexpr int THREADS_PER_WAVEFRONT = 32;
 
public:
    static constexpr const char* MMA_STORAGE_TYPE = "CUDA_MMA";
 
    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;
 
    data_flow::ImplementationType get_mma_impl_type(const GpuMmaTiling& tiling) const override;
 
    bool is_valid_block_size(
        const MmaBlockSize& block_size, types::PrimitiveType input_type, types::PrimitiveType acc_type
    ) const override;
 
    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;
 
    void set_mma_fragment_storage_type(
        types::StorageType& storage_type, const GpuMmaTiling& size, MmaFragmentType type, MmaFragmentLayout layout
    ) const override;
 
    std::vector<MmaBlockSize>
    get_supported_block_sizes(types::PrimitiveType input_type, types::PrimitiveType acc_type) const override;
 
    void emit_block_frag_type(
        std::ostream& os,
        codegen::LanguageExtension& lang_ext,
        const types::StorageType& storage_type,
        types::PrimitiveType element_type
    ) const;
 
    static void emit_block_frag_type(
        std::ostream& os,
        codegen::LanguageExtension& lang_ext,
        MmaFragmentType type,
        const MmaBlockSize& dims,
        MmaFragmentLayout layout,
        types::PrimitiveType scalar_type,
        std::optional<std::pair<int, int>> coop_dims
    );
 
    static bool is_mma_type(const types::StorageType& storage);
};
 
 
class CudaArch : public GpuArch {
    int sm_version_;
    CudaMmaSupport mma_support_;
 
public:
    CudaArch(int sm_version, bool mma_base_support, bool mma_tf32_support, bool mma_fp64_support)
        : GpuArch(), sm_version_(sm_version), mma_support_(mma_base_support, mma_tf32_support, mma_fp64_support) {
    }
 
    std::string unique_id() const override;
    std::string name() const override;
 
    int per_cu_threads() const override {
        return 32;
    }
 
    const CudaMmaSupport* 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 CudaArch CUDA_ARCH_SM70;
extern CudaArch CUDA_ARCH_SM75;
extern CudaArch CUDA_ARCH_SM80;
extern CudaArch CUDA_ARCH_SM86;
extern CudaArch CUDA_ARCH_SM87;
extern CudaArch CUDA_ARCH_SM89;
extern CudaArch CUDA_ARCH_SM90;
extern CudaArch CUDA_ARCH_SM100;
extern CudaArch CUDA_ARCH_SM103;
extern CudaArch CUDA_ARCH_SM107;
extern CudaArch CUDA_ARCH_SM110;
extern CudaArch CUDA_ARCH_SM120;

const CudaArch* cuda_arch_from_schedule_type(const structured_control_flow::ScheduleType& schedule);
const CudaArch* cuda_arch_from_env();
const CudaArch* cuda_arch_from_available_hardware();
const CudaArch* cuda_arch_parse(const std::string& name);
 
} // namespace sdfg::gpu::cuda