Arm backend: Remove duplicate scratch allocation in Direct Drive

backends/arm/runtime/EthosUBackend.cpp

@@ -15,7 +15,6 @@
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
-#include <memory>
 #include <new>
 #include <string>
 #include <vector>
@@ -162,20 +161,24 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface {
     const int input_count = handles.inputs ? handles.inputs->count : 0;
     const int output_count = handles.outputs ? handles.outputs->count : 0;
 
-    MemoryAllocator* temp_allocator = context.get_temp_allocator();
-    // Use a temporary allocator for the intermediate tensors of the
-    // computation. The allocator is released in runtime/executor/method.cpp at
-    // the end of the execution of the Ethos-U custom delegate
-    // Ethos-U driver requires 16 bit alignment.
-    char* ethosu_scratch = static_cast<char*>(
-        temp_allocator->allocate(handles.scratch_data_size, 16UL));
-    if (ethosu_scratch == nullptr) {
-      ET_LOG(
-          Error,
-          "Failed to allocate scratch buffer of %zu bytes from temp_allocator",
-          handles.scratch_data_size);
-      return Error::MemoryAllocationFailed;
+    char* ethosu_scratch = nullptr;
+    if (needs_scratch_allocation()) {
+      MemoryAllocator* temp_allocator = context.get_temp_allocator();
+      // Use a temporary allocator for the intermediate tensors of the
+      // computation. The allocator is released in runtime/executor/method.cpp
+      // at the end of the execution of the Ethos-U custom delegate. Ethos-U
+      // driver requires 16 bit alignment.
+      ethosu_scratch = static_cast<char*>(
+          temp_allocator->allocate(handles.scratch_data_size, 16UL));
+      if (ethosu_scratch == nullptr) {
+        ET_LOG(
+            Error,
+            "Failed to allocate scratch buffer of %zu bytes from temp_allocator",
+            handles.scratch_data_size);
+        return Error::MemoryAllocationFailed;
+      }
     }
+
     ET_LOG(
         Debug,
         "Running program data:\n  cmd %p %zu\n  weight %p %zu\n  scratch %p %zu\n  fast scratch %p %zu\n",
@@ -194,7 +197,6 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface {
     for (int i = 0; i < input_count; i++) {
       auto tensor_count = 1, io_count = 1;
       auto tensor_in = args[i]->toTensor();
-      char* scratch_addr = ethosu_scratch + handles.inputs->io[i].offset;
 
       // We accept:
       bool supported = 0;
@@ -228,30 +230,34 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface {
         return Error::InvalidProgram;
       }
 
-      // Select a compatible copy routine including checking for input layouts
-      // which require permutation.
-      bool both_int = tensor_in.scalar_type() == ScalarType::Int &&
-          handles.inputs->io[i].elem_size == 4;
-      bool both_char = (tensor_in.scalar_type() == ScalarType::Char ||
-                        tensor_in.scalar_type() == ScalarType::Byte) &&
-          handles.inputs->io[i].elem_size == 1;
-      bool both_short = tensor_in.scalar_type() == ScalarType::Short &&
-          handles.inputs->io[i].elem_size == 2;
-      bool both_bool = tensor_in.scalar_type() == ScalarType::Bool &&
-          (handles.inputs->io[i].elem_size == 1);
-
-      if (both_char || both_int || both_short || both_bool) {
-        EXECUTORCH_PROF_SCOPE(
-            event_tracer, "+EthosUBackend::execute()handles.input.memcpy()");
-        // Sizes match and elt size matches so memcpy.
-        // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate.
-        arm_ethos_io_memcpy(
-            scratch_addr,
-            tensor_in.mutable_data_ptr<char>(),
-            tensor_in.nbytes());
-      } else {
-        ET_LOG(Error, "No matching input copy routine");
-        return Error::InvalidProgram;
+      if (needs_scratch_allocation()) {
+        char* scratch_addr = ethosu_scratch + handles.inputs->io[i].offset;
+
+        // Select a compatible copy routine including checking for input layouts
+        // which require permutation.
+        bool both_int = tensor_in.scalar_type() == ScalarType::Int &&
+            handles.inputs->io[i].elem_size == 4;
+        bool both_char = (tensor_in.scalar_type() == ScalarType::Char ||
+                          tensor_in.scalar_type() == ScalarType::Byte) &&
+            handles.inputs->io[i].elem_size == 1;
+        bool both_short = tensor_in.scalar_type() == ScalarType::Short &&
+            handles.inputs->io[i].elem_size == 2;
+        bool both_bool = tensor_in.scalar_type() == ScalarType::Bool &&
+            (handles.inputs->io[i].elem_size == 1);
+
+        if (both_char || both_int || both_short || both_bool) {
+          EXECUTORCH_PROF_SCOPE(
+              event_tracer, "+EthosUBackend::execute()handles.input.memcpy()");
+          // Sizes match and elt size matches so memcpy.
+          // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate.
+          arm_ethos_io_memcpy(
+              scratch_addr,
+              tensor_in.mutable_data_ptr<char>(),
+              tensor_in.nbytes());
+        } else {
+          ET_LOG(Error, "No matching input copy routine");
+          return Error::InvalidProgram;
+        }
       }
       calculate_dimensions(
           tensor_in, &handles.inputs->io[i], &tensor_count, &io_count);

backends/arm/runtime/EthosUBackend_Cortex_A.cpp

@@ -186,7 +186,6 @@ Error invoke_linux_driver(
     ET_LOG(Error, "Ethos-U backend missing output metadata");
     return Error::InvalidProgram;
   }
-
   try {
     EthosU::Device& device = get_linux_device_cache().get(options.device_path);
     auto network = std::make_shared<EthosU::Network>(
@@ -343,8 +342,10 @@ void platform_destroy(PlatformState* state) {
   delete state;
 }
 
-// Used by EthosUBackend.cpp through EthosUBackend_Internal.h.
-// cppcheck-suppress unusedFunction
+bool needs_scratch_allocation() {
+  return false;
+}
+
 Error platform_execute(
     BackendExecutionContext& /*context*/,
     const ExecutionHandle* execution_handle,
@@ -394,7 +395,6 @@ Error platform_execute(
     ET_LOG(Error, "Ethos-U Linux backend missing platform state");
     return Error::InvalidState;
   }
-
   Error status = invoke_linux_driver(
       handles,
       linux_input_ptrs,

backends/arm/runtime/EthosUBackend_Cortex_M.cpp

@@ -64,6 +64,10 @@ void platform_destroy(PlatformState* state) {
   delete state;
 }
 
+bool needs_scratch_allocation() {
+  return true;
+}
+
 Error platform_execute(
     BackendExecutionContext& /*context*/,
     const ExecutionHandle* /*execution_handle*/,
@@ -72,6 +76,11 @@ Error platform_execute(
     int output_count,
     Span<executorch::runtime::EValue*> args,
     char* ethosu_scratch) {
+  if (handles.scratch_data_size > 0 && ethosu_scratch == nullptr) {
+    ET_LOG(Error, "Ethos-U scratch buffer is missing");
+    return Error::InvalidState;
+  }
+
   // Parse product config from command stream to reserve the correct driver
   uint32_t product, log2_macs;
   // The weak fallback below always returns 0, but some builds replace it

backends/arm/runtime/EthosUBackend_Internal.h

@@ -15,13 +15,12 @@
 #pragma GCC diagnostic ignored "-Wdouble-promotion"
 #endif
 
-#include <cstddef>
-#include <cstdint>
-
 #include <executorch/backends/arm/runtime/VelaBinStream.h>
 #include <executorch/runtime/backend/interface.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/evalue.h>
+#include <cstddef>
+#include <cstdint>
 
 #if defined(__GNUC__) && defined(__ZEPHYR__)
 #pragma GCC diagnostic pop
@@ -82,7 +81,11 @@ extern size_t ethosu_fast_scratch_size;
 PlatformState* platform_init(
     executorch::runtime::ArrayRef<executorch::runtime::CompileSpec> specs,
     executorch::runtime::MemoryAllocator* allocator);
+
 void platform_destroy(PlatformState* state);
+
+bool needs_scratch_allocation();
+
 executorch::runtime::Error platform_execute(
     executorch::runtime::BackendExecutionContext& context,
     const ExecutionHandle* execution_handle,