From 1b098d27d3f7974e8cbfec44aa1954f58f7635d4 Mon Sep 17 00:00:00 2001 From: Raman Sarokin Date: Fri, 30 Oct 2020 11:58:30 -0700 Subject: [PATCH] Added Buffer to Metal backend. PiperOrigin-RevId: 339913277 Change-Id: If4a92b3fafe922b5abf00bc8dc04deef1ef6ca6d --- tensorflow/lite/delegates/gpu/metal/BUILD | 29 ++++++ tensorflow/lite/delegates/gpu/metal/buffer.h | 95 +++++++++++++++++++ tensorflow/lite/delegates/gpu/metal/buffer.mm | 69 ++++++++++++++ .../lite/delegates/gpu/metal/buffer_test.mm | 70 ++++++++++++++ 4 files changed, 263 insertions(+) create mode 100644 tensorflow/lite/delegates/gpu/metal/buffer.h create mode 100644 tensorflow/lite/delegates/gpu/metal/buffer.mm create mode 100644 tensorflow/lite/delegates/gpu/metal/buffer_test.mm diff --git a/tensorflow/lite/delegates/gpu/metal/BUILD b/tensorflow/lite/delegates/gpu/metal/BUILD index cfefe53abb9..83aa333cdb2 100644 --- a/tensorflow/lite/delegates/gpu/metal/BUILD +++ b/tensorflow/lite/delegates/gpu/metal/BUILD @@ -48,6 +48,32 @@ cc_library( ], ) +objc_library( + name = "buffer", + srcs = ["buffer.mm"], + hdrs = ["buffer.h"], + copts = DEFAULT_COPTS, + sdk_frameworks = ["Metal"], + deps = [ + "//tensorflow/lite/delegates/gpu/common:status", + "@com_google_absl//absl/types:span", + ], +) + +objc_library( + name = "buffer_test_lib", + testonly = 1, + srcs = ["buffer_test.mm"], + sdk_frameworks = [ + "XCTest", + "Metal", + ], + deps = [ + ":buffer", + "//tensorflow/lite/delegates/gpu/common:types", + ], +) + objc_library( name = "buffer_convert", srcs = ["buffer_convert.mm"], @@ -285,6 +311,7 @@ objc_library( name = "common_tests_lib", testonly = 1, srcs = [ + "//tensorflow/lite/delegates/gpu/metal:buffer_test.mm", "//tensorflow/lite/delegates/gpu/metal:common_test.mm", "//tensorflow/lite/delegates/gpu/metal:compiled_model_test.mm", "//tensorflow/lite/delegates/gpu/metal:inference_context_test.mm", @@ -293,6 +320,8 @@ objc_library( ], sdk_frameworks = ["XCTest"], deps = [ + "//tensorflow/lite/delegates/gpu/common:types", + "//tensorflow/lite/delegates/gpu/metal:buffer", "//tensorflow/lite/delegates/gpu/metal:common", "//tensorflow/lite/delegates/gpu/metal:environment", "//tensorflow/lite/delegates/gpu/metal:inference_context", diff --git a/tensorflow/lite/delegates/gpu/metal/buffer.h b/tensorflow/lite/delegates/gpu/metal/buffer.h new file mode 100644 index 00000000000..65e54ca9225 --- /dev/null +++ b/tensorflow/lite/delegates/gpu/metal/buffer.h @@ -0,0 +1,95 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the Licensgoe is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#ifndef TENSORFLOW_LITE_DELEGATES_GPU_METAL_BUFFER_H_ +#define TENSORFLOW_LITE_DELEGATES_GPU_METAL_BUFFER_H_ + +#include +#include + +#import + +#include "absl/types/span.h" +#include "tensorflow/lite/delegates/gpu/common/status.h" + +namespace tflite { +namespace gpu { +namespace metal { + +class Buffer { + public: + Buffer() {} // just for using Buffer as a class members + Buffer(id buffer, size_t size_in_bytes); + + // Move only + Buffer(Buffer&& buffer); + Buffer& operator=(Buffer&& buffer); + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + + ~Buffer(); + + // for profiling and memory statistics + uint64_t GetMemorySizeInBytes() const { return size_; } + + id GetMemoryPtr() const { return buffer_; } + + // Writes data to a buffer. Data should point to a region that + // has exact size in bytes as size_in_bytes(constructor parameter). + template + absl::Status WriteData(const absl::Span data); + + // Reads data from Buffer into CPU memory. + template + absl::Status ReadData(std::vector* result) const; + + private: + void Release(); + + id buffer_ = nullptr; + size_t size_; +}; + +absl::Status CreateBuffer(size_t size_in_bytes, const void* data, id device, + Buffer* result); + +template +absl::Status Buffer::WriteData(const absl::Span data) { + if (size_ != sizeof(T) * data.size()) { + return absl::InvalidArgumentError( + "absl::Span data size is different from buffer allocated size."); + } + std::memcpy([buffer_ contents], data.data(), size_); + return absl::OkStatus(); +} + +template +absl::Status Buffer::ReadData(std::vector* result) const { + if (size_ % sizeof(T) != 0) { + return absl::UnknownError("Wrong element size(typename T is not correct?"); + } + + const int elements_count = size_ / sizeof(T); + result->resize(elements_count); + std::memcpy(result->data(), [buffer_ contents], size_); + + return absl::OkStatus(); +} + +} // namespace metal +} // namespace gpu +} // namespace tflite + +#endif // TENSORFLOW_LITE_DELEGATES_GPU_METAL_BUFFER_H_ diff --git a/tensorflow/lite/delegates/gpu/metal/buffer.mm b/tensorflow/lite/delegates/gpu/metal/buffer.mm new file mode 100644 index 00000000000..fa87faeb545 --- /dev/null +++ b/tensorflow/lite/delegates/gpu/metal/buffer.mm @@ -0,0 +1,69 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/delegates/gpu/metal/buffer.h" + +#include + +namespace tflite { +namespace gpu { +namespace metal { + +Buffer::Buffer(id buffer, size_t size_in_bytes) + : buffer_(buffer), size_(size_in_bytes) {} + +Buffer::Buffer(Buffer&& buffer) : buffer_(buffer.buffer_), size_(buffer.size_) { + buffer.buffer_ = nullptr; + buffer.size_ = 0; +} + +Buffer& Buffer::operator=(Buffer&& buffer) { + if (this != &buffer) { + Release(); + std::swap(size_, buffer.size_); + std::swap(buffer_, buffer.buffer_); + } + return *this; +} + +Buffer::~Buffer() { Release(); } + +void Buffer::Release() { + if (buffer_) { + buffer_ = nullptr; + size_ = 0; + } +} + +absl::Status CreateBuffer(size_t size_in_bytes, const void* data, + id device, Buffer* result) { + id buffer; + if (data) { + buffer = [device newBufferWithBytes:data + length:size_in_bytes + options:MTLResourceStorageModeShared]; + } else { + buffer = [device newBufferWithLength:size_in_bytes + options:MTLResourceStorageModeShared]; + } + + *result = Buffer(buffer, size_in_bytes); + + return absl::OkStatus(); +} + +} // namespace metal +} // namespace gpu +} // namespace tflite diff --git a/tensorflow/lite/delegates/gpu/metal/buffer_test.mm b/tensorflow/lite/delegates/gpu/metal/buffer_test.mm new file mode 100644 index 00000000000..70a63eb27ac --- /dev/null +++ b/tensorflow/lite/delegates/gpu/metal/buffer_test.mm @@ -0,0 +1,70 @@ +/* Copyright 2020 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "tensorflow/lite/delegates/gpu/metal/buffer.h" + +#include "tensorflow/lite/delegates/gpu/common/types.h" + +#import + +#import + +#include +#include + +@interface BufferTest : XCTestCase +@end + +@implementation BufferTest +- (void)setUp { + [super setUp]; +} + +using tflite::gpu::half; + +- (void)testBufferF32 { + id device = MTLCreateSystemDefaultDevice(); + + const std::vector data = {1.0f, 2.0f, 3.0f, -4.0f, 5.1f}; + tflite::gpu::metal::Buffer buffer; + XCTAssertTrue(tflite::gpu::metal::CreateBuffer(sizeof(float) * 5, nullptr, device, &buffer).ok()); + XCTAssertTrue(buffer.WriteData(absl::MakeConstSpan(data.data(), data.size())).ok()); + std::vector gpu_data; + XCTAssertTrue(buffer.ReadData(&gpu_data).ok()); + + XCTAssertEqual(gpu_data.size(), data.size()); + for (int i = 0; i < gpu_data.size(); ++i) { + XCTAssertEqual(gpu_data[i], data[i]); + } +} + +- (void)testBufferF16 { + id device = MTLCreateSystemDefaultDevice(); + + const std::vector data = {half(1.0f), half(2.0f), half(3.0f), half(-4.0f), half(5.1f)}; + tflite::gpu::metal::Buffer buffer; + XCTAssertTrue(tflite::gpu::metal::CreateBuffer( + sizeof(tflite::gpu::half) * 5, nullptr, device, &buffer).ok()); + XCTAssertTrue(buffer.WriteData(absl::MakeConstSpan(data.data(), data.size())).ok()); + std::vector gpu_data; + XCTAssertTrue(buffer.ReadData(&gpu_data).ok()); + + XCTAssertEqual(gpu_data.size(), data.size()); + for (int i = 0; i < gpu_data.size(); ++i) { + XCTAssertEqual(gpu_data[i], data[i]); + } +} + +@end