From 8fe56f080703c2ac6ae0b12b98d1c42e18ae85f2 Mon Sep 17 00:00:00 2001
From: Raman Sarokin <sorokin@google.com>
Date: Wed, 9 Dec 2020 10:46:16 -0800
Subject: [PATCH] TFLGpuDelegateSetCommandEncoder replaced by
 TFLGpuDelegateSetCommandBuffer.

PiperOrigin-RevId: 346586647
Change-Id: Ie40d5a9659cd4d031e8073f16fb78e31e7885a84
---
 tensorflow/lite/delegates/gpu/metal_delegate.mm         | 8 ++++++++
 tensorflow/lite/delegates/gpu/metal_delegate_internal.h | 5 +++++
 2 files changed, 13 insertions(+)

diff --git a/tensorflow/lite/delegates/gpu/metal_delegate.mm b/tensorflow/lite/delegates/gpu/metal_delegate.mm
index 60f8f361b35..d3eb35df407 100644
--- a/tensorflow/lite/delegates/gpu/metal_delegate.mm
+++ b/tensorflow/lite/delegates/gpu/metal_delegate.mm
@@ -716,6 +716,14 @@ bool TFLGpuDelegateBindMetalBufferToTensor(TfLiteDelegate* delegate, int tensor_
 
 // Note: This function is not exposed in `metal_delegate.h`, but it's exposed in
 // `metal_delegate_internal.h`.
+bool TFLGpuDelegateSetCommandEncoder(
+    TfLiteDelegate* delegate, id<MTLComputeCommandEncoder> encoder) {
+  auto* metal_delegate = ::tflite::gpu::metal::GetMetalDelegate(delegate);
+  if (!metal_delegate) return false;
+  metal_delegate->SetCommandEncoder(encoder);
+  return true;
+}
+
 bool TFLGpuDelegateSetCommandBuffer(TfLiteDelegate* delegate,
                                     id<MTLCommandBuffer> command_buffer) {
   auto* metal_delegate = ::tflite::gpu::metal::GetMetalDelegate(delegate);
diff --git a/tensorflow/lite/delegates/gpu/metal_delegate_internal.h b/tensorflow/lite/delegates/gpu/metal_delegate_internal.h
index 121caef450d..1a33d046103 100644
--- a/tensorflow/lite/delegates/gpu/metal_delegate_internal.h
+++ b/tensorflow/lite/delegates/gpu/metal_delegate_internal.h
@@ -33,6 +33,11 @@ bool TFLGpuDelegateBindMetalBufferToTensor(TfLiteDelegate* delegate,
                                            int tensor_index,
                                            id<MTLBuffer> metal_buffer);
 
+// Binds user-defined MTLComputeCommandEncoder. The delegate puts all GPU tasks
+// into this encoder instead of the internal encoder.
+bool TFLGpuDelegateSetCommandEncoder(TfLiteDelegate* delegate,
+                                     id<MTLComputeCommandEncoder> encoder);
+
 // Binds user-defined MTLCommandBuffer. The delegate puts all GPU tasks
 // into this buffer instead of the internal command buffer.
 bool TFLGpuDelegateSetCommandBuffer(TfLiteDelegate* delegate,