blob: 5e77df7d4add91984653f21a1abfe4cdf3ac14ed [file] [log] [blame]
/*
* Copyright (c) 2018-2020, 2022-2023 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 "kernel_weight_format.hpp"
#include <cstdint>
#include <functional>
namespace arm_gemm {
/* Structure describing an implementation. For each supported combination
* of types, a static list of these structures is built up to describe the
* implementations available.
*/
template<typename Top, typename Tret, class OutputStage = Nothing>
struct GemmImplementation {
const GemmMethod method;
const char * name;
const KernelWeightFormat kernel_weight_format = KernelWeightFormat::NON_FIXED;
std::function<bool(const GemmArgs &, const OutputStage &)> is_supported = {};
std::function<uint64_t(const GemmArgs &, const OutputStage &)> cycle_estimate = {};
std::function<GemmCommon<Top, Tret> *(const GemmArgs &, const OutputStage &)> instantiate = {};
bool do_is_supported(const GemmArgs &args, const OutputStage &os) const {
// Check supplied is_supported() function first.
if (is_supported != nullptr && !is_supported(args, os)) {
return false;
}
// Check weight format is appropriate.
if (args._fixed_format == false) {
// Can't return a fixed format kernel if we weren't asked for one.
return (kernel_weight_format == KernelWeightFormat::NON_FIXED);
} else {
// Fixed format kernel requested: if this is a non-fixed format kernel we can't use it.
if (kernel_weight_format == KernelWeightFormat::NON_FIXED) {
return false;
}
// If there's no config, or the config says ANY then this one is OK.
if (!args._cfg || args._cfg->weight_format == WeightFormat::ANY) {
return true;
}
// If we get here it means there is a config and it specifies a format. Check it matches this kernel.
// NOTE: this will execute SVE instructions if it's an SVE kernel, so it's important that is_supported()
// was called above first.
return (args._cfg->weight_format == get_weight_format(kernel_weight_format, sizeof(Top)));
}
}
uint64_t do_cycle_estimate(const GemmArgs &args, const OutputStage &os) const {
if (cycle_estimate != nullptr) {
return cycle_estimate(args, os);
} else {
return 0;
}
}
GemmCommon<Top, Tret> *do_instantiate(const GemmArgs &args, const OutputStage &os) const {
return instantiate(args, os);
}
static GemmImplementation with_estimate(GemmMethod m, const char *n,
std::function<bool(const GemmArgs &, const OutputStage &)> is_supported, std::function<uint64_t(const GemmArgs &, const OutputStage &)> cycle_estimate,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &, const OutputStage &)> instantiate) {
GemmImplementation impl(m,n);
impl.is_supported=is_supported;
impl.cycle_estimate=cycle_estimate;
impl.instantiate=instantiate;
return impl;
}
GemmImplementation(const GemmImplementation &) = default;
GemmImplementation & operator= (const GemmImplementation &) = default;
GemmImplementation(GemmMethod m, const char * n) : method(m), name(n) {}
GemmImplementation(GemmMethod m, const char *n,
std::function<bool(const GemmArgs &, const OutputStage &)> is_supported, std::function<bool(const GemmArgs &, const OutputStage &)> is_recommended,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &, const OutputStage &)> instantiate) :
method(m), name(n), is_supported(is_supported),
cycle_estimate( [is_recommended](const GemmArgs &args, const OutputStage &os) { return (is_recommended == nullptr) ? 0 : (is_recommended(args, os) ? 0 : UINT64_MAX); } ),
instantiate(instantiate) { }
GemmImplementation(GemmMethod m, const char *n, KernelWeightFormat kwf,
std::function<bool(const GemmArgs &, const OutputStage &)> is_supported, std::function<bool(const GemmArgs &, const OutputStage &)> is_recommended,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &, const OutputStage &)> instantiate) :
method(m), name(n), kernel_weight_format(kwf), is_supported(is_supported),
cycle_estimate( [is_recommended](const GemmArgs &args, const OutputStage &os) { return (is_recommended == nullptr) ? 0 : (is_recommended(args, os) ? 0 : UINT64_MAX); } ),
instantiate(instantiate) { }
};
/* Slightly different version of above for straightforward GEMMs with no
* output stage, so the std::functions there don't have to deal with the
* unnecessary second argument. */
template<typename Top, typename Tret>
struct GemmImplementation<Top, Tret, Nothing> {
const GemmMethod method;
const char * name;
const KernelWeightFormat kernel_weight_format = KernelWeightFormat::NON_FIXED;
std::function<bool(const GemmArgs &)> is_supported = {};
std::function<uint64_t(const GemmArgs &)> cycle_estimate = {};
std::function<GemmCommon<Top, Tret> *(const GemmArgs &)> instantiate = {};
bool do_is_supported(const GemmArgs &args, const Nothing &) const {
// Check supplied is_supported() function first.
if (is_supported != nullptr && !is_supported(args)) {
return false;
}
// Check weight format is appropriate.
if (args._fixed_format == false) {
// Can't return a fixed format kernel if we weren't asked for one.
return (kernel_weight_format == KernelWeightFormat::NON_FIXED);
} else {
// Fixed format kernel requested: if this is a non-fixed format kernel we can't use it.
if (kernel_weight_format == KernelWeightFormat::NON_FIXED) {
return false;
}
// If there's no config, or the config says ANY then this one is OK.
if (!args._cfg || args._cfg->weight_format == WeightFormat::ANY) {
return true;
}
// If we get here it means there is a config and it specifies a format. Check it matches this kernel.
// NOTE: this will execute SVE instructions if it's an SVE kernel, so it's important that is_supported()
// was called above first.
return (args._cfg->weight_format == get_weight_format(kernel_weight_format, sizeof(Top)));
}
}
uint64_t do_cycle_estimate(const GemmArgs &args, const Nothing &) const {
if (cycle_estimate != nullptr) {
return cycle_estimate(args);
} else {
return 0;
}
}
GemmCommon<Top, Tret> *do_instantiate(const GemmArgs &args, const Nothing &) const {
return instantiate(args);
}
static GemmImplementation with_estimate(GemmMethod m, const char *n,
std::function<bool(const GemmArgs &)> is_supported, std::function<uint64_t(const GemmArgs &)> cycle_estimate,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &)> instantiate) {
GemmImplementation impl(m,n);
impl.is_supported=is_supported;
impl.cycle_estimate=cycle_estimate;
impl.instantiate=instantiate;
return impl;
}
static GemmImplementation with_estimate(GemmMethod m, const char *n, KernelWeightFormat f,
std::function<bool(const GemmArgs &)> is_supported, std::function<uint64_t(const GemmArgs &)> cycle_estimate,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &)> instantiate) {
GemmImplementation impl(m,n,f);
impl.is_supported=is_supported;
impl.cycle_estimate=cycle_estimate;
impl.instantiate=instantiate;
return impl;
}
GemmImplementation(const GemmImplementation &) = default;
GemmImplementation & operator= (const GemmImplementation &) = default;
GemmImplementation(GemmMethod m, const char *n, KernelWeightFormat f=KernelWeightFormat::NON_FIXED) : method(m), name(n), kernel_weight_format(f) {}
GemmImplementation(GemmMethod m, const char *n,
std::function<bool(const GemmArgs &)> is_supported, std::function<bool(const GemmArgs &)> is_recommended,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &)> instantiate) :
method(m), name(n), is_supported(is_supported),
cycle_estimate( [is_recommended](const GemmArgs &args) -> uint64_t { return (is_recommended == nullptr) ? 0 : (is_recommended(args) ? 0 : UINT64_MAX); } ),
instantiate(instantiate) { }
GemmImplementation(GemmMethod m, const char *n, KernelWeightFormat kwf,
std::function<bool(const GemmArgs &)> is_supported, std::function<bool(const GemmArgs &)> is_recommended,
std::function<GemmCommon<Top, Tret> *(const GemmArgs &)> instantiate) :
method(m), name(n), kernel_weight_format(kwf), is_supported(is_supported),
cycle_estimate( [is_recommended](const GemmArgs &args) -> uint64_t { return (is_recommended == nullptr) ? 0 : (is_recommended(args) ? 0 : UINT64_MAX); } ),
instantiate(instantiate) { }
};
/* Provides the list of implementation descriptors which is processed by the
* other functions.
*
* A specialised version is provided for each supported combination of types.
* The end of the list is indicated by a sentinel descriptor with
* method==GemmMethod::DEFAULT. */
template<typename Top, typename Tret, class OutputStage = Nothing>
const GemmImplementation<Top, Tret, OutputStage> *gemm_implementation_list();
/*
* Select a GEMM implementation for the given arguments.
*
* The logic here returns the method on the list which supports the
* requested problem parameters, matches the provided filters (method and/or
* name string match) and offers the lowest cycle estimate. A cycle
* estimate of '0' is treated as a special value, causing the corresponding
* method to be selected immediately.
*
* If no method supports the requested parameters and passes the filters,
* this function returns false and doesn't touch the provided pointer
* reference.
*/
template<typename Top, typename Tret, class OutputStage>
bool find_implementation(const GemmArgs &args, const OutputStage &os, const GemmImplementation<Top, Tret, OutputStage> * &impl) {
auto gemms = gemm_implementation_list<Top, Tret, OutputStage>();
const GemmConfig *cfg = args._cfg;
const GemmImplementation<Top, Tret, OutputStage> *saved_impl = nullptr;
uint64_t best_estimate = 0;
for (const GemmImplementation<Top, Tret, OutputStage> *i = gemms; i->method != GemmMethod::DEFAULT; i++) {
/* Skip if this implementation doesn't support these args. */
if (!i->do_is_supported(args, os)) {
continue;
}
/* Skip if a specific method is requested and this is a different one. */
if (cfg && cfg->method != GemmMethod::DEFAULT && i->method != cfg->method) {
continue;
}
/* Skip if a filter is to be applied and it doesn't match. */
if (cfg && cfg->filter != "" && !strstr(i->name, cfg->filter.c_str())) {
continue;
}
/* Test the cycle estimate */
uint64_t estimate = i->do_cycle_estimate(args, os);
/* Short circuit - if the estimate is zero, return this one immediately. */
if (estimate==0) {
impl=i;
return true;
}
/* Otherwise, remember this is our best so far if we don't yet have
* a valid candidate, or we beat the estimate. */
if ((saved_impl == nullptr) || (estimate < best_estimate)) {
saved_impl = i;
best_estimate = estimate;
}
}
/* Return whichever method gave the best estimate. */
if (saved_impl != nullptr) {
impl = saved_impl;
return true;
}
return false;
}
template<typename Top, typename Tret, class OutputStage>
std::vector<KernelDescription> get_compatible_kernels(const GemmArgs &args, const OutputStage &os) {
std::vector<KernelDescription> res;
/* Find out what the default implementation in so we can set the flag accordingly later. */
const GemmImplementation<Top, Tret, OutputStage> *default_impl;
find_implementation(args, os, default_impl);
auto gemms = gemm_implementation_list<Top, Tret, OutputStage>();
for (const GemmImplementation<Top, Tret, OutputStage> *i = gemms; i->method != GemmMethod::DEFAULT; i++) {
/* Check that this implementation supports the presented problem. */
if (!i->do_is_supported(args, os)) {
continue;
}
res.push_back(KernelDescription(i->method, i->name, i==default_impl, i->do_cycle_estimate(args, os)));
}
return res;
}
template<typename Top, typename Tret, class OutputStage>
bool has_opt_gemm(WeightFormat &wf, const GemmArgs &args, const OutputStage &os) {
const GemmImplementation<Top, Tret, OutputStage> *impl;
const bool success = find_implementation<Top, Tret, OutputStage>(args, os, impl);
if (success)
wf = UniqueGemmCommon<Top, Tret>(impl->do_instantiate(args, os))->get_config().weight_format;
return success;
}
template<typename Top, typename Tret, class OutputStage>
UniqueGemmCommon<Top, Tret> gemm(const GemmArgs &args, const OutputStage &os) {
const GemmImplementation<Top, Tret, OutputStage> *impl;
if (find_implementation<Top, Tret, OutputStage>(args, os, impl)) {
return UniqueGemmCommon<Top, Tret>(impl->do_instantiate(args, os));
}
return UniqueGemmCommon<Top, Tret>(nullptr);
}
template<typename Top, typename Tret, class OutputStage>
KernelDescription get_gemm_method(const GemmArgs &args, const OutputStage &os) {
const GemmImplementation<Top, Tret, OutputStage> *impl;
if (find_implementation<Top, Tret>(args, os, impl)) {
return KernelDescription(impl->method, impl->name);
}
/* This shouldn't happen - there should always be at least one valid implementation. */
return KernelDescription();
}
} // namespace arm_gemm