Make HALIDE_REGISTER_GENERATOR work with multiple template args

src/Generator.h

@@ -1468,7 +1468,7 @@ class GIOBase {
     friend class GeneratorParamInfo;
 
     mutable int array_size_;  // always 1 if is_array() == false.
-        // -1 if is_array() == true but unspecified.
+                              // -1 if is_array() == true but unspecified.
 
     const std::string name_;
     const IOKind kind_;
@@ -2210,7 +2210,7 @@ class GeneratorInput : public Internal::GeneratorInputImplBase<T> {
     }
 
     // Avoid ambiguity between Func-with-dim and int-with-default
-    //template <typename T2 = T, typename std::enable_if<std::is_same<TBase, Func>::value>::type * = nullptr>
+    // template <typename T2 = T, typename std::enable_if<std::is_same<TBase, Func>::value>::type * = nullptr>
     GeneratorInput(size_t array_size, const std::string &name, IntIfNonScalar d)
         : Super(array_size, name, d) {
     }
@@ -3803,7 +3803,8 @@ struct halide_global_ns;
     namespace GEN_REGISTRY_NAME##_ns {                                                                                              \
         std::unique_ptr<Halide::Internal::GeneratorBase> factory(const Halide::GeneratorContext &context);                          \
         std::unique_ptr<Halide::Internal::GeneratorBase> factory(const Halide::GeneratorContext &context) {                         \
-            return GEN_CLASS_NAME::create(context, #GEN_REGISTRY_NAME, #FULLY_QUALIFIED_STUB_NAME);                                 \
+            using GenType = std::remove_pointer<decltype(new GEN_CLASS_NAME)>::type; /* NOLINT(bugprone-macro-parentheses) */       \
+            return GenType::create(context, #GEN_REGISTRY_NAME, #FULLY_QUALIFIED_STUB_NAME);                                        \
         }                                                                                                                           \
     }                                                                                                                               \
     static auto reg_##GEN_REGISTRY_NAME = Halide::Internal::RegisterGenerator(#GEN_REGISTRY_NAME, GEN_REGISTRY_NAME##_ns::factory); \

test/generator/CMakeLists.txt

@@ -435,6 +435,10 @@ halide_define_aot_test(string_param PARAMS "rpn_expr=5 y * x +")
 # shuffler_generator.cpp
 halide_define_aot_test(shuffler)
 
+# templated_aottest.cpp
+# templated_generator.cpp
+halide_define_aot_test(templated)
+
 # tiled_blur_aottest.cpp
 # tiled_blur_generator.cpp
 halide_define_aot_test(tiled_blur EXTRA_LIBS blur2x2)

test/generator/templated_aottest.cpp

@@ -0,0 +1,33 @@
+#include "HalideBuffer.h"
+#include "HalideRuntime.h"
+
+#include <math.h>
+#include <stdio.h>
+
+#include "templated.h"
+
+using namespace Halide::Runtime;
+
+int main(int argc, char **argv) {
+    const int kSize = 1024;
+
+    Buffer<double> output(kSize, kSize);
+    Buffer<float> input(kSize, kSize);
+
+    input.fill(17.0f);
+
+    templated(input, output);
+
+    auto check = [&](float val, float input_val) {
+        if (val != input_val + 2) {
+            printf("Output value was %f instead of %f\n",
+                   val, input_val + 2);
+            exit(-1);
+        }
+    };
+
+    output.for_each_value(check, input);
+
+    printf("Success!\n");
+    return 0;
+}

test/generator/templated_generator.cpp

@@ -0,0 +1,35 @@
+#include "Halide.h"
+
+using namespace Halide;
+
+// Define a templated generator. Normally this is a bad idea, and your template
+// parameters (e.g. the type of the input) should be GeneratorParams
+// instead. Sometimes, however, it's more convenient to have the C++ type
+// available as a template parameter. Or maybe you want to template a Generator
+// on something not expressible as a GeneratorParam. Or maybe you have a
+// deficient build system and it's difficult to specify GeneratorParams in the
+// build (note that HALIDE_REGISTER_GENERATOR_ALIAS also exists for this
+// purpose).
+template<typename T1, typename T2>
+class Templated : public Generator<Templated<T1, T2>> {
+public:
+    // A major downside of templated generators is that because we use CRTP, you
+    // must manually import names of types from the base class. For Input and
+    // Output you can also just use these equivalent globally-scoped names:
+    GeneratorInput<Buffer<T1>> input{"input", 2};
+    GeneratorOutput<Buffer<T2>> output{"output", 2};
+
+    void generate() {
+        Var x, y;
+        output(x, y) = cast<T2>(input(x, y) + (T1)2);
+
+        // Again, due to CRTP, we must manually tell C++ how to look up a method
+        // from the base class using this->template.
+        output.vectorize(x, this->template natural_vector_size<T2>());
+    }
+};
+
+// To pass a comma-separated template parameter list to a macro, we must enclose
+// the type argument in parentheses.
+HALIDE_REGISTER_GENERATOR((Templated<float, double>), templated)
+HALIDE_REGISTER_GENERATOR((Templated<uint8_t, uint16_t>), templated_uint8)