Merge pull request #7 from KernelTuner/develop

Add support for templated types as parameters in pragma kernels

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/compiler.h

@@ -101,8 +101,9 @@ struct Compiler: ICompiler {
      */
     template<typename C>
     Compiler(C&& compiler) :
-        inner_(std::make_shared<typename std::decay<C>::type>(
-            std::forward<C>(compiler))) {}
+        inner_(
+            std::make_shared<typename std::decay<C>::type>(
+                std::forward<C>(compiler))) {}
 
     void compile_ptx(
         KernelDef def,

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;
 

include/kernel_launcher/utils.h

@@ -2,6 +2,7 @@
 #define KERNEL_LAUNCHER_UTILS_H
 
 #include <cuda_runtime_api.h>
+#include <stdint.h>
 
 #include <functional>
 #include <iosfwd>
@@ -11,7 +12,6 @@
 #include <type_traits>
 #include <typeindex>
 #include <vector>
-#include <stdint.h>
 
 namespace kernel_launcher {
 

src/arg.cpp

@@ -37,7 +37,7 @@ KernelArg::KernelArg(KernelArg&& that) noexcept : KernelArg() {
 
 KernelArg::~KernelArg() {
     if (is_scalar() && !is_inline_scalar(type_)) {
-        delete[](char*) data_.large_scalar;
+        delete[] (char*)data_.large_scalar;
     }
 }
 

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());
 }
 
@@ -168,6 +180,14 @@ KernelBuilder& KernelBuilder::buffer_size(ArgExpr arg, TypedExpr<size_t> len) {
         ArgsEval eval {args, fallback};
         size_t i = arg.get();
         size_t n = eval(len);
+
+        if (i >= args.size()) {
+            throw std::runtime_error(
+                "argument " + std::to_string(i)
+                + " is out of bounds for kernel that has only "
+                + std::to_string(args.size()) + " parameters");
+        }
+
         args[i] = args[i].to_array(n);
     });
 }
@@ -444,7 +464,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 +516,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
+}  // 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;

src/internal/directives.cpp

@@ -215,7 +215,7 @@ static Expr parse_prim(TokenStream& stream, const Context& ctx) {
 
 // NOLINTNEXTLINE(readability-function-cognitive-complexity)
 static Expr parse_binop(TokenStream& stream, const Context& ctx, int prec) {
-    // TODO: == != <= >= && || %
+    // TODO: %
     Expr lhs = parse_prim(stream, ctx);
 
     while (true) {
@@ -416,6 +416,20 @@ static void parse_tuning_key_directive(
     builder.tuning_key(std::move(key));
 }
 
+template<typename... Ts>
+static bool one_of(const std::string& value, Ts&... items) {
+    static constexpr size_t N = sizeof...(Ts);
+    const char* array[N] = {items...};
+
+    for (size_t i = 0; i < N; i++) {
+        if (value == array[i]) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 static void
 process_directive(TokenStream& stream, KernelBuilder& builder, Context& ctx) {
     while (!stream.next_if(TokenKind::DirectiveEnd)) {
@@ -424,25 +438,28 @@ process_directive(TokenStream& stream, KernelBuilder& builder, Context& ctx) {
 
         if (name == "tune") {
             parse_tune_directive(stream, builder, ctx);
-        } else if (name == "set") {
+        } else if (one_of(name, "set", "let")) {
             parse_set_directive(stream, ctx);
-        } else if (name == "buffers" || name == "buffer") {
+        } else if (one_of(name, "buffer", "buffers")) {
             parse_buffer_directive(stream, builder, ctx);
         } else if (name == "tuning_key") {
             parse_tuning_key_directive(stream, builder, ctx);
-        } else if (name == "grid_size" || name == "grid_dim") {
+        } else if (one_of(name, "grid_size", "grid_dim")) {
             auto l = parse_expr_list3(stream, ctx);
             builder.grid_size(l[0], l[1], l[2]);
-        } else if (name == "block_size" || name == "block_dim") {
+        } else if (one_of(name, "block_size", "block_dim")) {
             auto l = parse_expr_list3(stream, ctx);
             builder.block_size(l[0], l[1], l[2]);
-        } else if (name == "grid_divisor" || name == "grid_divisors") {
+        } else if (one_of(name, "grid_divisor", "grid_divisors")) {
             auto l = parse_expr_list3(stream, ctx);
             builder.grid_divisors(l[0], l[1], l[2]);
-        } else if (name == "problem_size" || name == "problem_dim") {
+        } else if (one_of(name, "problem_size", "problem_dim")) {
             auto l = parse_expr_list3(stream, ctx);
             builder.problem_size(l[0], l[1], l[2]);
-        } else if (name == "restriction" || name == "restrictions") {
+        } else if (
+            name == "if"
+            || one_of(name, "restriction", "restrictions", "restrict")
+            || one_of(name, "assertion", "assertions", "assert")) {
             for (const auto& expr : parse_expr_list(stream, ctx)) {
                 builder.restriction(expr);
             }

src/internal/parser.cpp

@@ -46,8 +46,18 @@ static std::vector<FunctionParam> parse_kernel_params(TokenStream& stream) {
         Token before_name = begin;
         Token name = stream.next();
         Token end = stream.peek();
+        int template_depth = 0;
+
+        while (template_depth > 0
+               || !(
+                   end.kind == TokenKind::Comma
+                   || end.kind == TokenKind::ParenR)) {
+            if (name.kind == TokenKind::AngleL) {
+                template_depth++;
+            } else if (name.kind == TokenKind::AngleR && template_depth > 0) {
+                template_depth--;
+            }
 
-        while (end.kind != TokenKind::Comma && end.kind != TokenKind::ParenR) {
             before_name = name;
             name = stream.next();
             end = stream.peek();

src/internal/tokens.cpp

@@ -87,8 +87,9 @@ static index_t advance_string(index_t i, const std::string& input) {
 TokenKind char2_to_kind(char a, char b) {
     if ((a == '=' && b == '=') || (a == '!' && b == '=')
         || (a == '<' && b == '=') || (a == '>' && b == '=')
-        || (a == '&' && b == '&') || (a == '|' && b == '|')
-        || (a == '<' && b == '<') || (a == '>' && b == '>')
+        || (a == '&' && b == '&')
+        || (a == '|' && b == '|')
+        //|| (a == '<' && b == '<') || (a == '>' && b == '>')
         || (a == ':' && b == ':')) {
         return TokenKind::Punct;
     }