[HVX] Fix state_var issue (halide#6894)

* fix HVX state_var issue * abort if host is nullptr
ardier · Mar 3, 2024 · b764dfb · b764dfb
1 parent b99b41a
commit b764dfb
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 7 deletions.
diff --git a/src/CodeGen_Internal.cpp b/src/CodeGen_Internal.cpp
@@ -128,6 +128,7 @@ bool function_takes_user_context(const std::string &name) {
         "halide_hexagon_initialize_kernels",
         "halide_hexagon_run",
         "halide_hexagon_device_release",
+        "halide_hexagon_get_module_state",
         "halide_hexagon_power_hvx_on",
         "halide_hexagon_power_hvx_on_mode",
         "halide_hexagon_power_hvx_on_perf",

diff --git a/src/HexagonOffload.cpp b/src/HexagonOffload.cpp
@@ -695,12 +695,6 @@ class InjectHexagonRpc : public IRMutator {
 
     Module &device_code;
 
-    Expr state_var(const std::string &name, Type type) {
-        return Let::make(name, state_var_ptr(name, type),
-                         Load::make(type_of<void *>(), name, 0,
-                                    Buffer<>(), Parameter(), const_true(), ModulusRemainder()));
-    }
-
     Expr state_var_ptr(const std::string &name, Type type) {
         Expr &buf = state_bufs[name];
         if (!buf.defined()) {
@@ -712,7 +706,7 @@ class InjectHexagonRpc : public IRMutator {
     }
 
     Expr module_state() {
-        return state_var("hexagon_module_state", type_of<void *>());
+        return Call::make(type_of<void *>(), "halide_hexagon_get_module_state", {state_var_ptr("hexagon_module_state", type_of<void *>())}, Call::Extern);
     }
 
     Expr module_state_ptr() {

diff --git a/src/runtime/HalideRuntimeHexagonHost.h b/src/runtime/HalideRuntimeHexagonHost.h
@@ -50,6 +50,9 @@ extern int halide_hexagon_detach_device_handle(void *user_context, struct halide
 extern void *halide_hexagon_get_device_handle(void *user_context, struct halide_buffer_t *buf);
 extern uint64_t halide_hexagon_get_device_size(void *user_context, struct halide_buffer_t *buf);
 
+/** Return a pointer to the module_state. */
+extern void *halide_hexagon_get_module_state(void *user_context, void **host);
+
 /** Power HVX on and off. Calling a Halide pipeline will do this
  * automatically on each pipeline invocation; however, it costs a
  * small but possibly significant amount of time for short running

diff --git a/src/runtime/hexagon_host.cpp b/src/runtime/hexagon_host.cpp
@@ -755,6 +755,11 @@ WEAK uint64_t halide_hexagon_get_device_size(void *user_context, struct halide_b
     return handle->size;
 }
 
+WEAK void *halide_hexagon_get_module_state(void *user_context, void **host) {
+    halide_abort_if_false(user_context, host != nullptr);
+    return host[0];
+}
+
 WEAK int halide_hexagon_device_and_host_malloc(void *user_context, struct halide_buffer_t *buf) {
     debug(user_context) << "halide_hexagon_device_and_host_malloc called.\n";
     int result = halide_hexagon_device_malloc(user_context, buf);

diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp
@@ -99,6 +99,7 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = {
     (void *)&halide_hexagon_device_release,
     (void *)&halide_hexagon_get_device_handle,
     (void *)&halide_hexagon_get_device_size,
+    (void *)&halide_hexagon_get_module_state,
     (void *)&halide_hexagon_initialize_kernels,
     (void *)&halide_hexagon_finalize_kernels,
     (void *)&halide_hexagon_power_hvx_off,