Add call to ynn_optimize_subgraph

This currently happens implicitly in `ynn_create_runtime`, but that will not be the case soon. (Calling it multiple times is harmless.)

PiperOrigin-RevId: 824754921
This commit is contained in:
A. Unique TensorFlower 2025-10-27 17:47:01 -07:00 committed by TensorFlower Gardener
parent 1fc47fae8e
commit 6e82d4d96b
2 changed files with 6 additions and 0 deletions

View File

@ -161,6 +161,7 @@ cc_library(
"//xla/backends/cpu/runtime/ynnpack:ynn_interop", "//xla/backends/cpu/runtime/ynnpack:ynn_interop",
"//xla/hlo/ir:hlo", "//xla/hlo/ir:hlo",
"//xla/stream_executor:device_memory", "//xla/stream_executor:device_memory",
"//xla/tsl/platform:errors",
"//xla/tsl/platform:logging", "//xla/tsl/platform:logging",
"//xla/tsl/platform:statusor", "//xla/tsl/platform:statusor",
"@XNNPACK//ynnpack", "@XNNPACK//ynnpack",

View File

@ -40,6 +40,7 @@ limitations under the License.
#include "xla/primitive_util.h" #include "xla/primitive_util.h"
#include "xla/shape.h" #include "xla/shape.h"
#include "xla/stream_executor/device_memory.h" #include "xla/stream_executor/device_memory.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/logging.h" #include "xla/tsl/platform/logging.h"
#include "xla/tsl/platform/statusor.h" #include "xla/tsl/platform/statusor.h"
#include "xla/util.h" #include "xla/util.h"
@ -285,6 +286,10 @@ static absl::StatusOr<YnnSubgraph> EmitYnnSubgraph(
} }
} }
ynn_status status = ynn_optimize_subgraph(
subgraph.get(), /*threadpool=*/nullptr, /*flags=*/0);
TF_RETURN_IF_ERROR(YnnStatusToStatus(status));
return subgraph; return subgraph;
} }