From b764dfbd36e4d3adb46ce35fecba1600a20cd1ee Mon Sep 17 00:00:00 2001
From: Alexander Root <32245479+rootjalex@users.noreply.github.com>
Date: Mon, 22 Aug 2022 12:08:34 -0400
Subject: [PATCH] [HVX] Fix state_var issue (#6894)

* fix HVX state_var issue

* abort if host is nullptr
---
 src/CodeGen_Internal.cpp               | 1 +
 src/HexagonOffload.cpp                 | 8 +-------
 src/runtime/HalideRuntimeHexagonHost.h | 3 +++
 src/runtime/hexagon_host.cpp           | 5 +++++
 src/runtime/runtime_api.cpp            | 1 +
 5 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/CodeGen_Internal.cpp b/src/CodeGen_Internal.cpp
index 61ddd8c3cbe9..48e6d6f369c2 100644
--- 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
index 1f1ce00b3f1e..1e6de70c2e9b 100644
--- 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
index ff0092ffd111..66a2b3262169 100644
--- 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
index 9652057b5b4f..9c5337f9dd12 100644
--- 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
index c06141dcf267..e3a360f013e6 100644
--- 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,