Skip to content

Commit

Permalink
refactor: move tflite::Span to independent header (#2638)
Browse files Browse the repository at this point in the history
Hoist the universally-useful tflite::Span out of
codegen/runtime/micro_codegen_context.h and into an independent header
usable from elsewhere in the project.

BUG=#2636
  • Loading branch information
rkuester authored Jul 24, 2024
1 parent 569ed29 commit d619ad8
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 16 deletions.
1 change: 1 addition & 0 deletions codegen/runtime/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -14,5 +14,6 @@ cc_library(
"//tensorflow/lite/micro:micro_context",
"//tensorflow/lite/micro:micro_graph",
"//tensorflow/lite/micro:micro_log",
"//tensorflow/lite/micro:span",
],
)
17 changes: 1 addition & 16 deletions codegen/runtime/micro_codegen_context.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,25 +19,10 @@ limitations under the License.
#include "tensorflow/lite/c/common.h"
#include "tensorflow/lite/micro/micro_context.h"
#include "tensorflow/lite/micro/micro_graph.h"
#include "tensorflow/lite/micro/span.h"

namespace tflite {

// A poor man's std::span, we should consider using the Pigweed span instead.
template <typename T>
class Span {
public:
constexpr Span(T* data, size_t size) noexcept : data_(data), size_(size) {}

constexpr T& operator[](size_t idx) const noexcept { return *(data_ + idx); }

constexpr T* data() const noexcept { return data_; }
constexpr size_t size() const noexcept { return size_; }

private:
T* data_;
size_t size_;
};

struct Subgraph {
Span<const size_t> inputs;
Span<const size_t> outputs;
Expand Down
6 changes: 6 additions & 0 deletions tensorflow/lite/micro/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,12 @@ cc_library(
],
)

cc_library(
name = "span",
hdrs = ["span.h"],
copts = micro_copts(),
)

cc_library(
name = "system_setup",
srcs = [
Expand Down
41 changes: 41 additions & 0 deletions tensorflow/lite/micro/span.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
/* Copyright 2024 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.
*/

#ifndef TENSORFLOW_LITE_MICRO_SPAN_H_
#define TENSORFLOW_LITE_MICRO_SPAN_H_

#include <cstddef>

namespace tflite {

// A poor man's std::span, we should consider using the Pigweed span instead.
template <typename T>
class Span {
public:
constexpr Span(T* data, size_t size) noexcept : data_(data), size_(size) {}

constexpr T& operator[](size_t idx) const noexcept { return *(data_ + idx); }

constexpr T* data() const noexcept { return data_; }
constexpr size_t size() const noexcept { return size_; }

private:
T* data_;
size_t size_;
};

} // end namespace tflite

#endif // TENSORFLOW_LITE_MICRO_SPAN_H_

0 comments on commit d619ad8

Please sign in to comment.