Add support restrictions containing kernel arguments

Author: stijnh

include/kernel_launcher/builder.h

@@ -2,6 +2,7 @@
 #define KERNEL_LAUNCHER_BUILDER_H
 
 #include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include "kernel_launcher/arg.h"
@@ -24,11 +25,13 @@ struct KernelInstance {
         CudaModule module,
         std::array<TypedExpr<uint32_t>, 3> block_size,
         std::array<TypedExpr<uint32_t>, 3> grid_size,
-        TypedExpr<uint32_t> shared_mem) :
+        TypedExpr<uint32_t> shared_mem = 0,
+        std::vector<TypedExpr<bool>> assertions = {}) :
         module_(std::move(module)),
         block_size_(std::move(block_size)),
         grid_size_(std::move(grid_size)),
-        shared_mem_(std::move(shared_mem)) {}
+        shared_mem_(std::move(shared_mem)),
+        assertions_(std::move(assertions)) {}
 
     void launch(
         cudaStream_t stream,
@@ -52,6 +55,7 @@ struct KernelInstance {
     std::array<TypedExpr<uint32_t>, 3> block_size_ = {1, 1, 1};
     std::array<TypedExpr<uint32_t>, 3> grid_size_ = {0, 0, 0};
     TypedExpr<uint32_t> shared_mem_ = 0;
+    std::vector<TypedExpr<bool>> assertions_;
 };
 
 struct KernelBuilderSerializerHack;

include/kernel_launcher/config.h

@@ -203,18 +203,33 @@ struct ConfigSpace {
      */
     void restriction(TypedExpr<bool> e);
 
+    /**
+     * Returns the restrictions added by `restriction()`
+     */
+    const std::vector<TypedExpr<bool>>& restrictions() const {
+        return restrictions_;
+    }
+
     /**
      * Returns the default configuration for this configuration space.
      */
     Config default_config() const;
 
     /**
-     *  Check if the given configuration is a valid member of this configuration
-     *  space. This method essentially checks three things:
+     *  Check if the given configuration is a member of this configuration
+     *  space. This method essentially checks two things:
      *
      *  * Does the configuration contain the correct parameters.
      *  * Do these parameter contain valid values.
-     *  * Does the configuration meet the restrictions.
+     *
+     *  However, it does _not_ check if the configuration satisfies the
+     *  restrictions.
+     */
+    bool contains(const Eval& config) const;
+
+    /**
+     *  Check if the given configuration is a valid member of this configuration
+     *  space and also meets the restrictions.
      */
     bool is_valid(const Eval& config) const;
 

src/builder.cpp

@@ -131,6 +131,18 @@ void KernelInstance::launch(
         }
     }
 
+    // We check for assertions now after printing the debug information. This
+    // allows one to check the debugging output to see what where the arguments
+    // provided to kernel that caused the assertion to fail.
+    for (const auto& assertion : assertions_) {
+        if (!eval(assertion)) {
+            std::stringstream ss;
+            ss << "failed to launch kernel `" << module_.function_name()
+               << "`, assertion failed: `" << assertion.to_string() << "`";
+            throw std::runtime_error(ss.str());
+        }
+    }
+
     module_.launch(stream, grid_size, block_size, smem, ptrs.data());
 }
 
@@ -444,7 +456,34 @@ KernelInstance KernelBuilder::compile(
     const std::vector<TypeInfo>& param_types,
     const ICompiler& compiler,
     CudaContextHandle ctx) const {
+    if (!contains(config)) {
+        std::stringstream ss;
+        ss << "invalid configuration: `" << config << "`";
+        throw std::runtime_error(ss.str());
+    }
+
     DeviceAttrEval eval = {ctx.device(), config};
+    std::vector<TypedExpr<bool>> assertions;
+
+    for (const auto& restriction : restrictions()) {
+        auto r = restriction.resolve(eval);
+
+        if (!r.is_constant()) {
+            // Any restriction that contain kernel arguments cannot be resolved
+            // now at this moment. We add these to the list of assertions
+            // that will be checked each time the kernel gets launched.
+            assertions.emplace_back(r);
+            continue;
+        }
+
+        if (!r.eval(eval)) {
+            std::stringstream ss;
+            ss << "configuration `" << config
+               << "` does not meet the following restriction: `"
+               << restriction.to_string() << "`";
+            throw std::runtime_error(ss.str());
+        }
+    }
 
     if (!is_valid(eval)) {
         std::stringstream ss;
@@ -469,7 +508,8 @@ KernelInstance KernelBuilder::compile(
         std::move(module),
         std::move(block_size),
         std::move(grid_size),
-        shared_mem};
+        std::move(shared_mem),
+        std::move(assertions)};
 }
 
 }  // namespace kernel_launcher

src/config.cpp

@@ -111,7 +111,7 @@ Config ConfigSpace::default_config() const {
     return config;
 }
 
-bool ConfigSpace::is_valid(const Eval& config) const {
+bool ConfigSpace::contains(const Eval& config) const {
     for (const auto& p : params_) {
         Value v;
 
@@ -124,6 +124,14 @@ bool ConfigSpace::is_valid(const Eval& config) const {
         }
     }
 
+    return true;
+}
+
+bool ConfigSpace::is_valid(const Eval& config) const {
+    if (!contains(config)) {
+        return false;
+    }
+
     for (const auto& r : restrictions_) {
         if (!config(r)) {
             return false;