/* * Copyright (c) 2017-2024 Arm Limited. * * SPDX-License-Identifier: MIT * * Permission is hereby granted, free of charge, to any person obtaining a copy * of this software and associated documentation files (the "Software"), to * deal in the Software without restriction, including without limitation the * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or * sell copies of the Software, and to permit persons to whom the Software is * furnished to do so, subject to the following conditions: * * The above copyright notice and this permission notice shall be included in all * copies or substantial portions of the Software. * * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE * SOFTWARE. */ #include "arm_gemm.hpp" #include "gemm_common.hpp" #include "gemm_hybrid.hpp" #include "gemm_hybrid_indirect.hpp" #include "gemm_implementation.hpp" #include "gemm_interleaved.hpp" #include "gemv_batched.hpp" #include "gemv_pretransposed.hpp" #include "kernels/a32_sgemm_8x6.hpp" #ifdef ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #include "kernels/a64_ffhybrid_fp32_mla_6x16.hpp" #include "kernels/a64_ffhybrid_fp32bf16fp32_mmla_4x24.hpp" #include "kernels/a64_ffhybrid_fp32bf16fp32_mmla_6x16.hpp" #include "kernels/a64_ffinterleaved_bf16fp32_mmla_8x12.hpp" #include "kernels/a64_ffinterleaved_fp32_mla_8x12.hpp" #endif // ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #include "kernels/a64_hybrid_fp32bf16fp32_mmla_4x24.hpp" #include "kernels/a64_hybrid_fp32bf16fp32_mmla_6x16.hpp" #include "kernels/a64_hybrid_fp32_mla_4x24.hpp" #include "kernels/a64_hybrid_fp32_mla_6x16.hpp" #include "kernels/a64_hybrid_fp32_mla_8x4.hpp" #include "kernels/a64_interleaved_bf16fp32_mmla_8x12.hpp" #include "kernels/a64_sgemm_8x12.hpp" #include "kernels/a64_sgemm_8x6.hpp" #include "kernels/a64_smallK_hybrid_fp32_mla_6x4.hpp" #include "kernels/a64_smallK_hybrid_fp32_mla_8x4.hpp" #ifdef ARM_COMPUTE_ENABLE_SVE #ifdef ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #include "kernels/sve_ffhybrid_fp32_mla_6x4VL.hpp" #include "kernels/sve_ffhybrid_fp32bf16fp32_mmla_4x6VL.hpp" #include "kernels/sve_ffinterleaved_fp32_mla_8x3VL.hpp" #include "kernels/sve_ffinterleaved_bf16fp32_mmla_8x3VL.hpp" #endif // ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #ifdef ARM_COMPUTE_ENABLE_SME2 #include "kernels/sme2_gemv_fp32_mla_16VL.hpp" #include "kernels/sme2_gemv_fp32bf16fp32_dot_16VL.hpp" #include "kernels/sme2_interleaved_nomerge_fp32_mopa_1VLx4VL.hpp" #include "kernels/sme2_interleaved_nomerge_bf16fp32_mopa_1VLx4VL.hpp" #include "kernels/sme2_interleaved_nomerge_fp32_mopa_2VLx2VL.hpp" #include "kernels/sme2_interleaved_nomerge_bf16fp32_mopa_2VLx2VL.hpp" #include "kernels/sme2_interleaved_nomerge_fp32_mopa_4VLx1VL.hpp" #include "kernels/sme2_interleaved_nomerge_bf16fp32_mopa_4VLx1VL.hpp" #endif // ARM_COMPUTE_ENABLE_SME2 #include "kernels/sve_ffhybrid_fp32_mla_6x4VL.hpp" #include "kernels/sve_ffhybrid_fp32bf16fp32_mmla_4x6VL.hpp" #include "kernels/sve_ffinterleaved_fp32_mla_8x3VL.hpp" #include "kernels/sve_ffinterleaved_bf16fp32_mmla_8x3VL.hpp" #include "kernels/sve_hybrid_fp32bf16fp32_mmla_4x6VL.hpp" #include "kernels/sve_hybrid_fp32bf16fp32_mmla_6x4VL.hpp" #include "kernels/sve_hybrid_fp32_mla_6x4VL.hpp" #include "kernels/sve_hybrid_fp32_mla_8x1VL.hpp" #include "kernels/sve_interleaved_bf16fp32_mmla_8x3VL.hpp" #include "kernels/sve_interleaved_fp32_mla_8x3VL.hpp" #include "kernels/sve_interleaved_fp32_mmla_8x3VL.hpp" #endif // ARM_COMPUTE_ENABLE_SVE namespace arm_gemm { static const GemmImplementation gemm_fp32_methods[] = { // GEMV cases - starting with 'gemv_batched' wrapper to turn batched GEMV into GEMM. { GemmMethod::GEMV_BATCHED, "gemv_batched", [](const GemmArgs &args) { return args._Msize==1 && args._nbatches>1 && !args._indirect_input; }, nullptr, [](const GemmArgs &args) { return new GemvBatched(args); } }, #ifdef __aarch64__ #ifdef ARM_COMPUTE_ENABLE_BF16 // "fast mode" (BF16) kernels GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "a64_interleaved_bf16fp32_mmla_8x12", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmInterleaved::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleaved(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_hybrid_fp32bf16fp32_mmla_6x16", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_hybrid_fp32bf16fp32_mmla_4x24", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), #endif // ARM_COMPUTE_ENABLE_BF16 #ifdef ARM_COMPUTE_ENABLE_SVE #ifdef ARM_COMPUTE_ENABLE_SME2 // SME kernels { GemmMethod::GEMM_HYBRID, "sme2_gemv_fp32bf16fp32_dot_16VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_sme2() && args._Msize==1 && args._nbatches==1 && !args._indirect_input; }, nullptr, [](const GemmArgs &args) { return new GemvPretransposed(args); } }, { GemmMethod::GEMM_HYBRID, "sme2_gemv_fp32_mla_16VL", [](const GemmArgs &args) { return args._ci->has_sme2() && args._Msize==1 && args._nbatches==1 && !args._indirect_input; }, nullptr, [](const GemmArgs &args) { return new GemvPretransposed(args); } }, #ifdef ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_bf16fp32_mopa_1VLx4VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_sme2(); }, [](const GemmArgs &args) { const auto VL = sme::get_vector_length(); return args._Msize <= VL || (2*VL < args._Msize && args._Msize <= 3*VL); }, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #endif // ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_fp32_mopa_1VLx4VL", [](const GemmArgs &args) { return args._ci->has_sme2(); }, [](const GemmArgs &args) { const auto VL = sme::get_vector_length(); return args._Msize <= VL || (2*VL < args._Msize && args._Msize <= 3*VL); }, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #ifdef ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_bf16fp32_mopa_4VLx1VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_sme2(); }, [](const GemmArgs &args) { const auto VL = sme::get_vector_length(); return args._Nsize <= VL || (2*VL < args._Nsize && args._Nsize <= 3*VL); }, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #endif // ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_fp32_mopa_4VLx1VL", [](const GemmArgs &args) { return args._ci->has_sme2(); }, [](const GemmArgs &args) { const auto VL = sme::get_vector_length(); return args._Nsize <= VL || (2*VL < args._Nsize && args._Nsize <= 3*VL); }, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #ifdef ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_bf16fp32_mopa_2VLx2VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_sme2(); }, nullptr, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #endif // ARM_COMPUTE_ENABLE_BF16 { GemmMethod::GEMM_INTERLEAVED, "sme2_interleaved_nomerge_fp32_mopa_2VLx2VL", [](const GemmArgs &args) { return args._ci->has_sme2(); }, nullptr, [](const GemmArgs &args) { return new GemmInterleavedNoMerge(args); } }, #endif // ARM_COMPUTE_ENABLE_SME2 #ifdef ARM_COMPUTE_ENABLE_BF16 GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "sve_interleaved_bf16fp32_mmla_8x3VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_svebf16(); }, [](const GemmArgs &args) { return GemmInterleaved::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleaved(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "sve_hybrid_fp32bf16fp32_mmla_6x4VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "sve_hybrid_fp32bf16fp32_mmla_4x6VL", [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), #endif // ARM_COMPUTE_ENABLE_BF16 #ifdef ARM_COMPUTE_ENABLE_SVEF32MM // MMLA next due to higher throughput (which is SVE only) // Prefer this in all cases, except if fast mode is requested and BF16 is available. { GemmMethod::GEMM_INTERLEAVED, "sve_interleaved_fp32_mmla_8x3VL", [](const GemmArgs &args) { return args._ci->has_svef32mm() && (args._Ksize>4); }, [](const GemmArgs &args) { return !(args._fast_mode && args._ci->has_bf16()); }, [](const GemmArgs &args) { return new GemmInterleaved(args); } }, #endif // ARM_COMPUTE_ENABLE_SVEF32MM // SVE kernels { GemmMethod::GEMM_HYBRID, "sve_hybrid_fp32_mla_8x1VL", [](const GemmArgs &args) { return args._ci->has_sve(); }, [](const GemmArgs &args) { return (args._Nsize < 12); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } }, GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "sve_hybrid_fp32_mla_6x4VL", [](const GemmArgs &args) { return args._ci->has_sve(); }, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "sve_interleaved_fp32_mla_8x3VL", [](const GemmArgs &args) { return args._ci->has_sve(); }, [](const GemmArgs &args) { return GemmInterleaved::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleaved(args); } ), #ifdef ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #ifdef ARM_COMPUTE_ENABLE_BF16 GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "sve_ffinterleaved_bf16fp32_mmla_8x3VL", KernelWeightFormat::VL2VL_BL64_BF16, [](const GemmArgs &args) { return args._fast_mode && args._ci->has_svebf16(); }, [](const GemmArgs &args) { return GemmInterleavedFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleavedFixedFormat(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "sve_ffhybrid_fp32bf16fp32_mmla_4x6VL", KernelWeightFormat::VL2VL_BL64_BF16, [](const GemmArgs &args) { return args._fast_mode && args._ci->has_svebf16(); }, [](const GemmArgs &args) { return GemmHybridIndirectFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirectFixedFormat(args); } ), #endif GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "sve_ffinterleaved_fp32_mla_8x3VL", KernelWeightFormat::VL1VL_BL32, [](const GemmArgs &args) { return args._ci->has_sve(); }, [](const GemmArgs &args) { return GemmInterleavedFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleavedFixedFormat(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "sve_ffhybrid_fp32_mla_6x4VL", KernelWeightFormat::VL1VL_BL32, [](const GemmArgs &args) { return args._ci->has_sve(); }, [](const GemmArgs &args) { return GemmHybridIndirectFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirectFixedFormat(args); } ), #endif // ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #endif // ARM_COMPUTE_ENABLE_SVE // Cortex-A35 specific kernel - use for any problem on A35, and never in any other cases. { GemmMethod::GEMM_INTERLEAVED, "a64_sgemm_8x6", nullptr, [](const GemmArgs &args) { return args._ci->get_cpu_model() == CPUModel::A35; }, [](const GemmArgs &args) { return new GemmInterleaved(args); } }, // Arm® Neon™ hybrid methods { GemmMethod::GEMM_HYBRID, "a64_smallK_hybrid_fp32_mla_8x4", [](const GemmArgs &args) { return args._Ksize <= 8 && (args._Nsize % 4)==0 && !args._indirect_input; }, nullptr, [](const GemmArgs &args) { return new GemmHybrid(args); } }, { GemmMethod::GEMM_HYBRID, "a64_smallK_hybrid_fp32_mla_6x4", [](const GemmArgs &args) { return (args._Ksize > 8 && args._Ksize <= 16) && (args._Nsize % 4)==0 && !args._indirect_input; }, nullptr, [](const GemmArgs &args) { return new GemmHybrid(args); } }, { GemmMethod::GEMM_HYBRID, "a64_hybrid_fp32_mla_8x4", nullptr, [](const GemmArgs &args) { return (args._Nsize < 12); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } }, GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_hybrid_fp32_mla_4x24", nullptr, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_hybrid_fp32_mla_6x16", nullptr, [](const GemmArgs &args) { return GemmHybridIndirect::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirect(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "a64_sgemm_8x12", nullptr, [](const GemmArgs &args) { return GemmInterleaved::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleaved(args); } ), #ifdef ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #ifdef ARM_COMPUTE_ENABLE_BF16 // "fast mode" (BF16) kernels GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "a64_ffinterleaved_bf16fp32_mmla_8x12", KernelWeightFormat::VL256_BL64_BF16, [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmInterleavedFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleavedFixedFormat(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_ffhybrid_fp32bf16fp32_mmla_4x24", KernelWeightFormat::VL256_BL64_BF16, [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirectFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirectFixedFormat(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_ffhybrid_fp32bf16fp32_mmla_6x16", KernelWeightFormat::VL256_BL64_BF16, [](const GemmArgs &args) { return args._fast_mode && args._ci->has_bf16(); }, [](const GemmArgs &args) { return GemmHybridIndirectFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirectFixedFormat(args); } ), #endif // BF16 GemmImplementation::with_estimate( GemmMethod::GEMM_INTERLEAVED, "a64_ffinterleaved_fp32_mla_8x12", KernelWeightFormat::VL128_BL32, nullptr, [](const GemmArgs &args) { return GemmInterleavedFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmInterleavedFixedFormat(args); } ), GemmImplementation::with_estimate( GemmMethod::GEMM_HYBRID, "a64_ffhybrid_fp32_mla_6x16", KernelWeightFormat::VL128_BL32, nullptr, [](const GemmArgs &args) { return GemmHybridIndirectFixedFormat::estimate_cycles(args); }, [](const GemmArgs &args) { return new GemmHybridIndirectFixedFormat(args); } ), #endif // ARM_COMPUTE_ENABLE_FIXED_FORMAT_KERNELS #endif // __aarch64__ #ifdef __arm__ { GemmMethod::GEMM_INTERLEAVED, "sgemm_8x6", nullptr, nullptr, [](const GemmArgs &args) { return new GemmInterleaved(args); } }, #endif // __arm__ { GemmMethod::DEFAULT, "", nullptr, nullptr, nullptr } }; /* Templated function to return this list. */ template<> const GemmImplementation *gemm_implementation_list() { return gemm_fp32_methods; } /* Explicitly instantiate the external functions for these types. */ template UniqueGemmCommon gemm(const GemmArgs &args, const Nothing &); template bool has_opt_gemm(WeightFormat &weight_format, const GemmArgs &args, const Nothing &); template KernelDescription get_gemm_method(const GemmArgs &args, const Nothing &); template std::vector get_compatible_kernels (const GemmArgs &args, const Nothing &); } // namespace arm_gemm