Add an ENVIRONMENT option to add_halide_library

cmake/HalideGeneratorHelpers.cmake

@@ -719,7 +719,7 @@ function(add_halide_library TARGET)
         FILE_BASE_NAME FROM GENERATOR FUNCTION_NAME NAMESPACE OUTPUT_DIR USE_RUNTIME AUTOSCHEDULER
         HEADER ${extra_output_names} NO_THREADS NO_DL_LIBS
     )
-    set(multiValueArgs DEPENDS TARGETS PARAMS PLUGINS ${features_args})
+    set(multiValueArgs DEPENDS ENVIRONMENT TARGETS PARAMS PLUGINS ${features_args})
     cmake_parse_arguments(ARG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
     if (NOT "${ARG_UNPARSED_ARGUMENTS}" STREQUAL "")
@@ -882,6 +882,17 @@ function(add_halide_library TARGET)
         OUT_DEPENDS generator_cmd_deps
     )
 
+    foreach (assignment IN LISTS ARG_ENVIRONMENT)
+        if (NOT assignment MATCHES "^[A-Za-z_][A-Za-z0-9_]*=")
+            message(
+                FATAL_ERROR "ENVIRONMENT entries must have the form VAR=value, got: ${assignment}"
+            )
+        endif ()
+    endforeach ()
+    if (ARG_ENVIRONMENT)
+        list(PREPEND generator_cmd ${CMAKE_COMMAND} -E env ${ARG_ENVIRONMENT} --)
+    endif ()
+
     set(generator_args
         COMMAND ${generator_cmd}  #
         DEPENDS ${generator_cmd_deps} ${ARG_DEPENDS}  #

doc/GeneratorCache.md

@@ -56,6 +56,7 @@ following are guaranteed to produce identical outputs:
   of its command-line params changed. Autoscheduling itself is deliberately
   excluded so that computing a cache key never requires running a (potentially
   expensive) autoscheduler.
+- **The value of `HL_LLVM_ARGS`**, since it passes options directly to LLVM.
 - **The contents of any `-p` plugins** (e.g. autoschedulers), since they affect
   codegen but live outside libHalide.
 

doc/HalideCMakePackage.md

@@ -414,6 +414,7 @@ add_halide_library(<target> FROM <generator-target>
                    [USE_RUNTIME hl-target]
                    [PARAMS param1 [param2 ...]]
                    [DEPENDS [dep1 dep2 ...]]
+                   [ENVIRONMENT VAR1=value1 [VAR2=value2 ...]]
                    [TARGETS target1 [target2 ...]]
                    [FEATURES feature1 [feature2 ...]]
                    [FEATURES[<triple>] feature1 [feature2 ...]]
@@ -491,6 +492,11 @@ listed here will be passed verbatim to `add_custom_command`, and so will be
 relative to the source directory. Use absolute paths when referring to files
 outside the source directory.
 
+`ENVIRONMENT` sets environment variables (each `VAR=value`) for the generator
+invocations of this library only, via `cmake -E env`. For example,
+`ENVIRONMENT HL_LLVM_ARGS=-misched-regpressure=false` passes a flag to LLVM for
+just this library's code generation. It does not affect the runtime target.
+
 To use an autoscheduler, set the `AUTOSCHEDULER` argument to a target named like
 `Namespace::Scheduler`, for example `Halide::Adams2019`. This will set the
 `autoscheduler` GeneratorParam on the generator command line to `Scheduler` and

src/Generator.cpp

@@ -1013,6 +1013,12 @@ void run_with_cache(const std::string &key,
 }
 
 #ifdef WITH_SERIALIZATION
+// HL_LLVM_ARGS is passed straight to LLVM's option parser, so it can change
+// codegen without changing anything else in the key.
+void add_llvm_args_key_field(CacheKeyBuilder &b) {
+    b.add(get_env_variable("HL_LLVM_ARGS"));
+}
+
 // Mix the invocation-wide (target-independent) inputs that affect codegen
 // into a fresh key builder: the compiler identity plus names, build mode,
 // requested output types, and the canonical generator-param settings.
@@ -1025,6 +1031,7 @@ bool add_common_key_fields(CacheKeyBuilder &b, const char *kind, const ExecuteGe
     }
     b.add(kind);
     b.add(compiler_id);
+    add_llvm_args_key_field(b);
     b.add(args.function_name);
     b.add(args.file_base_name);
     b.add(args.build_mode == ExecuteGeneratorArgs::Gradient ? "gradient" : "default");
@@ -1178,6 +1185,7 @@ void execute_generator(const ExecuteGeneratorArgs &args_in) {
                 CacheKeyBuilder b;
                 b.add("runtime");
                 b.add(compiler_id);
+                add_llvm_args_key_field(b);
                 b.add(gcd_target.to_string());
                 for (OutputFileType t : args.output_types) {
                     b.add(std::to_string((int)t));

test/correctness/generator_cache.cpp

@@ -90,14 +90,18 @@ std::vector<uint8_t> read_all(const fs::path &p) {
     return {std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
 }
 
-void set_cache_dir(const std::string &dir) {
+void set_env(const char *name, const std::string &value) {
 #ifdef _WIN32
-    _putenv_s("HL_CACHE_DIR", dir.c_str());
+    _putenv_s(name, value.c_str());
 #else
-    setenv("HL_CACHE_DIR", dir.c_str(), /*overwrite*/ 1);
+    setenv(name, value.c_str(), /*overwrite*/ 1);
 #endif
 }
 
+void set_cache_dir(const std::string &dir) {
+    set_env("HL_CACHE_DIR", dir);
+}
+
 // Find the single cached blob for a given output type (blobs are named "f<int>"
 // where <int> is the OutputFileType value). Returns an empty path if not found.
 fs::path find_blob(const fs::path &entries, OutputFileType type) {
@@ -169,6 +173,14 @@ int main(int argc, char **argv) {
     const fs::path obj_b = run("b", "1");
     check(read_all(obj_b) == sentinel);
 
+    // HL_LLVM_ARGS can change codegen, so it is part of the key: setting it
+    // misses instead of restoring the sentinel.
+    set_env("HL_LLVM_ARGS", "-misched-regpressure=false");
+    const fs::path obj_llvm_args = run("llvm-args", "1");
+    set_env("HL_LLVM_ARGS", "");
+    check(fs::exists(obj_llvm_args));
+    check(read_all(obj_llvm_args) != sentinel);
+
     // 3. A different GeneratorParam changes the pipeline, so the key differs and
     //    the run recompiles rather than restoring the sentinel.
     const fs::path obj_c = run("c", "2");

test/integration/CMakeLists.txt

@@ -169,6 +169,25 @@ add_test(
         --test-command ${CMAKE_CTEST_COMMAND} --output-on-failure
 )
 
+# ENVIRONMENT must reach the generator process: an unknown LLVM option passed
+# through HL_LLVM_ARGS makes the generator fail.
+add_test(
+    NAME aot_generator_environment
+    COMMAND
+        ${CMAKE_CTEST_COMMAND}
+        --build-and-test "${CMAKE_CURRENT_LIST_DIR}/aot" "${CMAKE_CURRENT_BINARY_DIR}/aot-environment"
+        --build-generator Ninja
+        --build-options -DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER} -DCMAKE_BUILD_TYPE=Release
+        "-Daot_ENVIRONMENT=HL_LLVM_ARGS=--unknown-option-from-halide-integration-test"
+        --test-command ${CMAKE_CTEST_COMMAND} --output-on-failure
+)
+set_tests_properties(
+    aot_generator_environment
+    PROPERTIES
+    PASS_REGULAR_EXPRESSION
+        "Unknown command line argument '--unknown-option-from-halide-integration-test'"
+)
+
 ##
 # Generator compile cache (HL_CACHE_DIR) tests
 #

test/integration/aot/CMakeLists.txt

@@ -19,6 +19,11 @@ elseif (aot_USE_PLUGINS_DIRECTLY)
     set(extra_options PLUGINS Halide::Adams2019 PARAMS autoscheduler=Adams2019)
 endif ()
 
+set(aot_ENVIRONMENT "" CACHE STRING "Environment assignments for the generator invocation")
+if (aot_ENVIRONMENT)
+    list(APPEND extra_options ENVIRONMENT ${aot_ENVIRONMENT})
+endif ()
+
 add_executable(add_gen add.cpp)
 target_link_libraries(add_gen PRIVATE Halide::Generator)