Merge pull request #32 from flagos-ai/fix/string-constexpr-arg

Support string constexpr kernel args; scope googletest to examples

Author: tengqm

CMakeLists.txt

@@ -223,13 +223,15 @@ add_subdirectory(src)
 # ==============================================================================
 # Examples
 # ==============================================================================
-set(INSTALL_GTEST OFF)
-github_url(GTEST_URL "google/googletest/archive/refs/tags/v1.16.0.tar.gz")
-FetchContent_Declare(googletest URL ${GTEST_URL} DOWNLOAD_EXTRACT_TIMESTAMP ON)
-FetchContent_MakeAvailable(googletest)
-
 option(TRITON_JIT_BUILD_EXAMPLES "Build examples" ${PROJECT_IS_TOP_LEVEL})
 if(TRITON_JIT_BUILD_EXAMPLES)
+    # googletest is only needed by the examples; fetch it here so that
+    # subdirectory consumers (e.g. FlagGems via FetchContent) do not pull it in.
+    set(INSTALL_GTEST OFF)
+    github_url(GTEST_URL "google/googletest/archive/refs/tags/v1.16.0.tar.gz")
+    FetchContent_Declare(googletest URL ${GTEST_URL} DOWNLOAD_EXTRACT_TIMESTAMP ON)
+    FetchContent_MakeAvailable(googletest)
+
     add_subdirectory(examples)
 endif()
 

include/triton_jit/triton_jit_function.h

@@ -5,6 +5,7 @@
 #include <optional>
 #include <sstream>
 #include <string>
+#include <string_view>
 #include <type_traits>
 #include <unordered_map>
 #include <utility>
@@ -147,6 +148,16 @@ struct ArgHandle {
       // Assumption: nullopt is always treated as constexpr,
       // even if the parameter is not marked as constexpr
       signature.push_back("nullopt");
+    } else if constexpr (std::is_same_v<std::decay_t<T>, const char*> ||
+                         std::is_same_v<std::decay_t<T>, char*> ||
+                         is_same_ignore_cvref<std::string, T>::value ||
+                         is_same_ignore_cvref<std::string_view, T>::value) {
+      // A string argument (e.g. a dtype spelled "tl.float32") can only ever be
+      // a constexpr Triton parameter. Route it to handle_constexpr at compile
+      // time so the specialized / non-constexpr paths -- which require
+      // triton_type<T>::name and do not specialize for strings -- are never
+      // instantiated for a string type.
+      handle_constexpr(item);
     } else {
       if (ssig.at(idx) == ArgType::CONSTEXPR) {  // constexpr
         handle_constexpr(item);