Support for priority levels of kernels in the catalog.

- Situation:
  - So far, the DAPHNE compiler selected the kernel for a DaphneIR operation from the kernel catalog based on argument/result types and backend.
  - (Besides that, the user can specify a kernel hint in DaphneDSL, which will be respected by the compiler.)
  - In case there are multiple applicable kernels, the DAPHNE compiler simply used the first one it finds.
- This commit introduces a priority to kernels registered in the kernel catalog.
  - The priority is a signed integer (0 by default).
  - If multiple kernels are applicable, the DAPHNE compiler chooses the one with the highest priority (e.g., a priority of 1 is better than 0); if there are multiple kernels with the highest priority, any of them is used.
  - All built-in kernels (those mentioned in kernels.json) are registered with the default priority of 0.
  - When registering custom kernel extensions, the user can optionally specify a priority on the command line.
  - A meaningful use of this feature is to make the DAPHNE compiler use the kernels provided by a custom extension whenever they are applicable (prefer them over the built-in kernels).
    - So far, the alternative would have been using manual kernel hints in DaphneDSL, but this approch can be cumbersome/error-prone and is not possible when the custom kernel implements a DaphneIR operation that has to direct counterpart in DaphneDSL.
  - Extended the test cases for kernel extensions.
  - Updated the documentation on extending DAPHNE (also added it to the navigation, where it was missing so far).
- Other noteworthy changes:
  - Errors happening during the population of the kernel catalog at DAPHNE start-up result in a parser error now (with the respective status code).

Author: pdamme

doc/Extensions.md

@@ -173,9 +173,12 @@ make
 ### Step 3: Using a Kernel Extension
 
 The kernels in a kernel extension can be used either automatically by DAPHNE or manually by the user.
-Automatic use is currently restricted to the selection of the kernel based on result/argument data/value types, but in the future we plan to support custom cost models as well.
-Besides that, the manual employment of custom kernels is very useful for experimentation, e.g., to see the impact of the kernel in the context of a larger integrated data analysis pipeline.
-To this end, DaphneDSL [compiler hints](/doc/DaphneDSL/LanguageRef.md#compiler-hints) tell DAPHNE to use a specific kernel, even though DAPHNE's optimizing compiler may not choose the kernel, otherwise.
+The manual use has precedence over the automatic use.
+
+#### Manual Use of Custom Kernels
+
+The manual employment of custom kernels is very useful for experimentation, e.g., to see the impact of a particular kernel at a certain point of a larger integrated data analysis pipeline.
+To this end, DaphneDSL [compiler hints](/doc/DaphneDSL/LanguageRef.md#compiler-hints) tell DAPHNE to use a specific kernel in a specific place, even though DAPHNE's optimizing compiler may not choose the kernel, otherwise.
 
 *Running example:*
 
@@ -205,7 +208,7 @@ s = sum::mySumSeq(X);
 print(s);
 ```
 
-We execute this script with the following command, whereby the argument `--kernel-ext` specified the kernel catalog JSON file of the extension to use:
+We execute this script with the following command, whereby the argument `--kernel-ext` specifies the kernel catalog JSON file of the extension to use:
 ```bash
 bin/daphne --kernel-ext scripts/examples/extensions/myKernels/myKernels.json scripts/examples/extensions/myKernels/demoSeq.daphne
 ```
@@ -222,4 +225,27 @@ print(s);
 We execute this script by:
 ```bash
 bin/daphne --kernel-ext scripts/examples/extensions/myKernels/myKernels.json scripts/examples/extensions/myKernels/demoSIMD.daphne
+```
+
+#### Automatic Use of Custom Kernels
+
+The automatic use of custom kernels is currently restricted to the selection of a kernel based on its result/argument data/value types and its priority level.
+In the future we plan to support custom cost models as well.
+
+*Running example:*
+
+Continuing the running example from above, we can make DAPHNE use the custom kernels `mySumSeq()` or `mySumSIMD()` even without a manual kernel hint by specifying a suitable *priority* when registering the `myKernels` extension with DAPHNE.
+
+Priority levels can optionally be specified with the `--kernel-ext` command line argument by appending a colon (`:`) followed by the priority as an integer.
+The default priority of `0` is used for all built-in kernels and for extension kernels in case no priority is specified.
+When registering a kernel extension, the given priority is assigned to *all* kernels provided by the extension.
+When multiple kernels are applicable for an operation based on the combination of argument/result data/value types as well as the backend, DAPHNE chooses the kernel with the highest priority.
+If there are multiple kernels with the highest priority, it is not specified which of them is used.
+
+By registering a kernel extension with a priority greater than zero, one can enforce that the kernels provided by the extension are always preferred over the built-in ones whenever they are applicable.
+For instance, the following command registers the `myKernels` extension with a priority of `1`.
+As the `myKernels` extension provides two kernels for the same operation, argument/result types, and backend, we cannot tell, based on priorities, which of these kernels will be used, but we can be sure that the built-in kernel will not be employed.
+
+```bash
+bin/daphne --kernel-ext scripts/examples/extensions/myKernels/myKernels.json:1 scripts/examples/extensions/myKernels/demo.daphne
 ```

mkdocs.yml

@@ -65,6 +65,7 @@ nav:
     - FPGAconfiguration.md
     - MPI-Usage.md
     - Profiling.md
+    - 'Custom Extensions': Extensions.md
   - 'Developers':
     - development/Contributing.md
     - development/BuildingDaphne.md

src/api/internal/daphne_internal.cpp

@@ -571,14 +571,48 @@ int startDAPHNE(int argc, const char **argv, DaphneLibResult *daphneLibRes, int
     // ************************************************************************
 
     KernelCatalog &kc = executor.getUserConfig().kernelCatalog;
-    // kc.dump();
-    KernelCatalogParser kcp(mctx);
-    kcp.parseKernelCatalog(user_config.libdir + "/catalog.json", kc);
-    if (user_config.use_cuda)
-        kcp.parseKernelCatalog(user_config.libdir + "/CUDAcatalog.json", kc);
-    // kc.dump();
-    if (!kernelExt.empty())
-        kcp.parseKernelCatalog(kernelExt, kc);
+    try {
+        // kc.dump();
+        KernelCatalogParser kcp(mctx);
+        kcp.parseKernelCatalog(user_config.libdir + "/catalog.json", kc, 0);
+        if (user_config.use_cuda)
+            kcp.parseKernelCatalog(user_config.libdir + "/CUDAcatalog.json", kc, 0);
+        // kc.dump();
+        if (!kernelExt.empty()) {
+            std::string extCatalogFile;
+            int64_t extPriority;
+
+            const std::string prioritySep = ":";
+            const size_t pos = kernelExt.rfind(prioritySep);
+            if (pos != std::string::npos) { // a priority was specified for the extension
+                extCatalogFile = kernelExt.substr(0, pos);
+                const std::string extPriorityStr(kernelExt.substr(pos + prioritySep.size()));
+                try {
+                    size_t idx;
+                    extPriority = std::stoll(extPriorityStr, &idx);
+                    if (idx != extPriorityStr.size())
+                        // stoll() did not consume all characters in extPriorityStr, there is some non-integer part at
+                        // the end of the string.
+                        throw std::runtime_error(""); // the error message is generated in the catch-block below
+                } catch (std::exception &e) {
+                    throw std::runtime_error("invalid priority for kernel extension, expected an integer after the '" +
+                                             prioritySep + "', but found '" + extPriorityStr + "': '" + kernelExt +
+                                             "'");
+                }
+            } else { // no priority was specified for the extension
+                extCatalogFile = kernelExt;
+                extPriority = 0;
+            }
+
+            kcp.parseKernelCatalog(extCatalogFile, kc, extPriority);
+        }
+    } catch (std::exception &e) {
+        logErrorDaphneLibAware(daphneLibRes, "Parser error: " + std::string(e.what()));
+        return StatusCode::PARSER_ERROR;
+    } catch (...) {
+        logErrorDaphneLibAware(daphneLibRes, "Parser error: Unknown exception");
+        return StatusCode::PARSER_ERROR;
+    }
 
     // ************************************************************************
     // Parse, compile and execute DaphneDSL script

src/compiler/catalog/KernelCatalog.h

@@ -60,9 +60,18 @@ struct KernelInfo {
      */
     const std::string libPath;
 
+    /**
+     * @brief The priority of this kernel.
+     *
+     * If there are multiple applicable kernels, the one with the highest priority will be preferred.
+     */
+    const int64_t priority;
+
     KernelInfo(const std::string kernelFuncName, const std::vector<mlir::Type> resTypes,
-               const std::vector<mlir::Type> argTypes, const std::string backend, const std::string libPath)
-        : kernelFuncName(kernelFuncName), resTypes(resTypes), argTypes(argTypes), backend(backend), libPath(libPath) {
+               const std::vector<mlir::Type> argTypes, const std::string backend, const std::string libPath,
+               int64_t priority)
+        : kernelFuncName(kernelFuncName), resTypes(resTypes), argTypes(argTypes), backend(backend), libPath(libPath),
+          priority(priority) {
         //
     }
 };

src/compiler/lowering/RewriteToCallKernelOpPass.cpp

@@ -35,6 +35,7 @@
 #include "llvm/ADT/ArrayRef.h"
 
 #include <iostream>
+#include <limits>
 #include <memory>
 #include <sstream>
 #include <stdexcept>
@@ -431,7 +432,8 @@ class KernelReplacement : public RewritePattern {
             const size_t numArgs = lookupArgTys.size();
             const size_t numRess = lookupResTys.size();
             int chosenKernelIdx = -1;
-            for (size_t i = 0; i < kernelInfos.size() && chosenKernelIdx == -1; i++) {
+            int64_t chosenKernelPriority = std::numeric_limits<int64_t>::min();
+            for (size_t i = 0; i < kernelInfos.size(); i++) {
                 auto ki = kernelInfos[i];
                 if (ki.backend != backend)
                     continue;
@@ -447,8 +449,11 @@ class KernelReplacement : public RewritePattern {
                 for (size_t i = 0; i < numRess && !mismatch; i++)
                     if (lookupResTys[i] != ki.resTypes[i])
                         mismatch = true;
-                if (!mismatch)
+
+                if (!mismatch && (ki.priority > chosenKernelPriority || chosenKernelIdx == -1)) {
                     chosenKernelIdx = i;
+                    chosenKernelPriority = ki.priority;
+                }
             }
             if (chosenKernelIdx == -1) {
                 std::stringstream s;

src/parser/catalog/KernelCatalogParser.cpp

@@ -109,7 +109,7 @@ void KernelCatalogParser::mapTypes(const std::vector<std::string> &in, std::vect
     }
 }
 
-void KernelCatalogParser::parseKernelCatalog(const std::string &filePath, KernelCatalog &kc) const {
+void KernelCatalogParser::parseKernelCatalog(const std::string &filePath, KernelCatalog &kc, int64_t priority) const {
     std::filesystem::path dirPath = std::filesystem::path(filePath).parent_path();
     try {
         std::ifstream kernelsConfigFile(filePath);
@@ -130,7 +130,7 @@ void KernelCatalogParser::parseKernelCatalog(const std::string &filePath, Kernel
             mapTypes(kernelData["resTypes"], resTypes, "result", kernelFuncName, opMnemonic, backend);
             std::vector<mlir::Type> argTypes;
             mapTypes(kernelData["argTypes"], argTypes, "argument", kernelFuncName, opMnemonic, backend);
-            kc.registerKernel(opMnemonic, KernelInfo(kernelFuncName, resTypes, argTypes, backend, libPath));
+            kc.registerKernel(opMnemonic, KernelInfo(kernelFuncName, resTypes, argTypes, backend, libPath, priority));
         }
     } catch (std::exception &e) {
         throw std::runtime_error("error while parsing kernel catalog file `" + filePath + "`: " + e.what());

src/parser/catalog/KernelCatalogParser.h

@@ -64,6 +64,7 @@ class KernelCatalogParser {
      *
      * @param filePath The path to the file to extract kernel information from.
      * @param kc The kernel catalog to register the kernels with.
+     * @param priority The kernel priority to use for all kernels parsed from the given file.
      */
-    void parseKernelCatalog(const std::string &filePath, KernelCatalog &kc) const;
+    void parseKernelCatalog(const std::string &filePath, KernelCatalog &kc, int64_t priority) const;
 };

src/runtime/distributed/worker/WorkerImpl.cpp

@@ -75,9 +75,9 @@ WorkerImpl::Status WorkerImpl::Compute(std::vector<WorkerImpl::StoredInfo> *outp
 
     KernelCatalog &kc = executor.getUserConfig().kernelCatalog;
     KernelCatalogParser kcp(executor.getContext());
-    kcp.parseKernelCatalog(cfg.libdir + "/catalog.json", kc);
+    kcp.parseKernelCatalog(cfg.libdir + "/catalog.json", kc, 0);
     if (executor.getUserConfig().use_cuda)
-        kcp.parseKernelCatalog(cfg.libdir + "/CUDAcatalog.json", kc);
+        kcp.parseKernelCatalog(cfg.libdir + "/CUDAcatalog.json", kc, 0);
 
     mlir::OwningOpRef<mlir::ModuleOp> module(mlir::parseSourceString<mlir::ModuleOp>(mlirCode, executor.getContext()));
     if (!module) {

test/api/cli/extensibility/ExtensionTest.cpp

@@ -52,13 +52,74 @@ TEST_CASE("extension_kernel", TAG_EXTENSIBILITY) {
     REQUIRE(status == 0);
 
     // *************************************************************************
-    // Use the custom kernel extension.
+    // Use the custom kernel extension based on hints or priority.
     // *************************************************************************
-    // Run a DaphneDSL script which uses a kernel from the extension through a
-    // kernel hint. The extension is registered with DAPHNE at run-time. DAPHNE
-    // itself is not re-built or anything.
-    compareDaphneToStr("hello from mySumAll\n2\n", std::string(dirPath + "extension_kernel_usage.daphne").c_str(),
-                       "--kernel-ext", std::string(dirPath + "kernel_extension_test/myKernels.json").c_str());
+    // All of the following invocations of DAPHNE register the extension with DAPHNE at run-time. DAPHNE itself is not
+    // re-built or anything. This extension provides a custom kernel for a case (DaphneIR operation, input/output types,
+    // backend) that is already covered by an existing built-in kernel.
+
+    // If a kernel hint is given, this kernel must always be used, irrespective of its priority.
+    // Run a DaphneDSL script which uses a kernel from the extension through a kernel hint.
+    SECTION("hint, no priority") { // must use the custom kernel
+        compareDaphneToStr("hello from mySumAll\n2\n",
+                           std::string(dirPath + "extension_kernel_usage_hint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json").c_str());
+    }
+    SECTION("hint, default priority") { // must use the custom kernel
+        compareDaphneToStr("hello from mySumAll\n2\n",
+                           std::string(dirPath + "extension_kernel_usage_hint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json:0").c_str());
+    }
+    SECTION("hint, higher priority") { // must use the custom kernel
+        compareDaphneToStr("hello from mySumAll\n2\n",
+                           std::string(dirPath + "extension_kernel_usage_hint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json:1").c_str());
+    }
+    SECTION("hint, lower priority") { // must use the custom kernel
+        compareDaphneToStr("hello from mySumAll\n2\n",
+                           std::string(dirPath + "extension_kernel_usage_hint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json:-1").c_str());
+    }
+    // If no kernel hint is given, the kernel must be used when it has a higher-than-default priority and must not be
+    // used if it has a lower-than-default priority. For "no hint, no priority" and "no hint, default priority",
+    // DAPHNE's choice whether to use the custom or the built-in kernel is unspecified, so we don't test these cases
+    // here.
+    // Run a DaphneDSL script which does not use a kernel hint.
+    SECTION("no hint, higher priority") { // must use the custom kernel
+        compareDaphneToStr("hello from mySumAll\n2\n",
+                           std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json:1").c_str());
+    }
+    SECTION("no hint, lower priority") { // must NOT use the custom kernel
+        compareDaphneToStr("2\n", std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                           std::string(dirPath + "kernel_extension_test/myKernels.json:-1").c_str());
+    }
+    // If an invalid value is specified for the priority, DAPHNE must stop, even if the kernel extension itself exists.
+    SECTION("no hint, invalid priority (empty)") {
+        checkDaphneStatusCode(StatusCode::PARSER_ERROR,
+                              std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                              std::string(dirPath + "kernel_extension_test/myKernels.json:").c_str());
+    }
+    SECTION("no hint, invalid priority (float)") {
+        checkDaphneStatusCode(StatusCode::PARSER_ERROR,
+                              std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                              std::string(dirPath + "kernel_extension_test/myKernels.json:0.1").c_str());
+    }
+    SECTION("no hint, invalid priority (string)") {
+        checkDaphneStatusCode(StatusCode::PARSER_ERROR,
+                              std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                              std::string(dirPath + "kernel_extension_test/myKernels.json:abc").c_str());
+    }
+    SECTION("no hint, invalid priority (integer and string)") {
+        checkDaphneStatusCode(StatusCode::PARSER_ERROR,
+                              std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(), "--kernel-ext",
+                              std::string(dirPath + "kernel_extension_test/myKernels.json:123abc").c_str());
+    }
+    SECTION("no hint, invalid priority (too huge integer)") {
+        checkDaphneStatusCode(
+            StatusCode::PARSER_ERROR, std::string(dirPath + "extension_kernel_usage_nohint.daphne").c_str(),
+            "--kernel-ext", std::string(dirPath + "kernel_extension_test/myKernels.json:99999999999999999999").c_str());
+    }
 
     // Clear the streams.
     out.clear();