wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 6 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/c/experimental/filesystem/BUILD
+77
View File
@@ -0,0 +1,77 @@
# Experimental filesystem C APIs for TensorFlow.
# Will be moved in proper place once all filesystems are converted to the
# modular framework.
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# This is only for plugins
cc_library(
name = "filesystem_interface",
hdrs = ["filesystem_interface.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status",
],
)
# Core TensorFlow depends on this, will be included in main library
cc_library(
name = "modular_filesystem",
srcs = [
"modular_filesystem.cc",
"modular_filesystem_registration.cc",
],
hdrs = [
"modular_filesystem.h",
"modular_filesystem_registration.h",
],
# TODO(b/139060984): Visibility should be more restrictive once we
# convert to modular filesystems everywhere
visibility = ["//visibility:public"],
deps = [
":filesystem_interface",
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:file_statistics",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:strcat",
"//tensorflow/core/platform:stringpiece",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
# Compliance test for modules and for interface
tf_cc_test(
name = "modular_filesystem_test",
size = "small",
srcs = ["modular_filesystem_test.cc"],
linkopts = ["-ldl"],
tags = [
"manual", # Requires DSOs as arguments, eventual setup
"notap", # b/139060984, requires implementing modular support for Google filesystem
],
deps = [
":modular_filesystem",
"//tensorflow/core:framework_internal",
"//tensorflow/core/lib/io:path",
"//tensorflow/core/platform:env",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:stacktrace_handler",
"//tensorflow/core/platform:test",
],
)
tensorflow/c/experimental/filesystem/filesystem_interface.h
+998
View File
@@ -0,0 +1,998 @@
/* Copyright 2019 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_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
#include <stddef.h>
#include <stdint.h>
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
/// This is the interop header between core TensorFlow and modular filesystem
/// plugins (see initial RFC https://github.com/tensorflow/community/pull/101).
///
/// Both core TensorFlow and every plugin will use this header. The associated
/// `.cc` file is only used by core TensorFlow to implement checking needed for
/// plugin registration and ensuring API and ABI compatibility. Plugin authors
/// don't need to read the `.cc` file but they should consult every section of
/// this file to ensure a compliant plugin can be built and that the plugin can
/// be used without recompilation in the widest range of TensorFlow versions.
///
/// The header is divided into sections, as follows:
/// 1. Opaque plugin private data structures and wrappers for type safety;
/// 2. Function tables for plugin functionality;
/// 3. Versioning metadata;
/// 4. Plugin registration API and the DSO entry point.
#ifdef __cplusplus
extern "C" {
#endif // __cplusplus
/// SECTION 1. Opaque data structures to hold plugin specific data
/// ----------------------------------------------------------------------------
///
/// The following data structures incorporate a `void*` that is opaque to
/// TensorFlow but can be used by each filesystem plugin to represent internal
/// data.
///
/// We prefer to have these structures instead of passing `void*` into
/// method signatures to have some type of type safety: for example, operations
/// that are only valid on random access files have a `TF_RandomAccessFile`
/// argument.
///
/// Lifetime: The wrapper data structures are owned by core TensorFlow. The data
/// pointed to by the `void*` members is always owned by the plugin. The plugin
/// will provide functions to call to allocate and deallocate this data (see
/// next sections) and core TensorFlow ensures to call these at the proper time.
///
/// Plugins will never receive a `TF_*` pointer that is `nullptr`. Core
/// TensorFlow will never touch the `void*` wrapped by these structures, except
/// to initialize it as `nullptr`.
typedef struct TF_RandomAccessFile {
void* plugin_file;
} TF_RandomAccessFile;
typedef struct TF_WritableFile {
void* plugin_file;
} TF_WritableFile;
typedef struct TF_ReadOnlyMemoryRegion {
void* plugin_memory_region;
} TF_ReadOnlyMemoryRegion;
typedef struct TF_Filesystem {
void* plugin_filesystem;
} TF_Filesystem;
// The named union is needed here (as opposed to
// inside the `TF_Filesystem_Option_Value` struct)
// as MSVC does not recognize `typeof`.
typedef union TF_Filesystem_Option_Value_Union {
int64_t int_val;
double real_val;
struct {
char* buf;
int buf_length;
} buffer_val;
} TF_Filesystem_Option_Value_Union;
typedef struct TF_Filesystem_Option_Value {
int type_tag; // type of values in the values union
int num_values; // number of values
TF_Filesystem_Option_Value_Union*
values; // owned (plugins must make a copy if storing this)
} TF_Filesystem_Option_Value;
typedef enum TF_Filesystem_Option_Type {
TF_Filesystem_Option_Type_Int = 0,
TF_Filesystem_Option_Type_Real,
TF_Filesystem_Option_Type_Buffer,
TF_Filesystem_Num_Option_Types, // must always be the last item
} TF_Filesystem_Option_Type;
typedef struct TF_Filesystem_Option {
char* name; // null terminated, owned
char* description; // null terminated, owned
int per_file; // bool actually, but bool is not a C type
TF_Filesystem_Option_Value* value; // owned
} TF_Filesystem_Option;
/// SECTION 2. Function tables for functionality provided by plugins
/// ----------------------------------------------------------------------------
///
/// The following data structures represent the function tables for operations
/// that plugins provide (some are mandatory, some are optional, with or without
/// a default implementation).
///
/// Each plugin implements the operations that are supported and TensorFlow will
/// properly handle the cases when an operation is not supported (i.e., return
/// the corresponding `Status` value).
///
/// REQUIRED OPERATIONS: All required operations are marked as such, including
/// operations which are conditionally required. If the presence of an operation
/// `foo` requires operation `bar` to be present, this is specified in `foo`. If
/// the entire set of operations in a table is not provided, use `nullptr` for
/// the struct pointer (e.g., when a file type is not supported).
///
/// DEFAULT IMPLEMENTATIONS: Some operations have default implementations that
/// TensorFlow uses in case the plugin doesn't supply its own version. An
/// operation `foo` might have a default implementation which uses `bar` and
/// `foobar`. If the plugin supplies `bar` and `foobar`, TensorFlow can use the
/// default implementation of `foo`.
///
/// During plugin loading, plugins will call the registration function provided
/// by this interface, supplying values for each of these structures. Core
/// TensorFlow checks that the plugin supplies all mandatory operations and
/// then copies these tables to a different memory location, marking the new
/// operation tables as read-only. Once a plugin is loaded, none of these
/// operation pointers may change.
///
/// There are 4 function tables: one for each of the 3 file objects in
/// TensorFlow (i.e., `RandomAccessFile`, `WritableFile`,
/// `ReadOnlyMemoryRegion`) and one for all the operations a `Filesystem`
/// implements. Each of them is in a 1-to-1 correspondence with the wrapper
/// structures from the first section: these tables only contain function
/// pointers that operate on the corresponding data. Thus, the first argument of
/// each of these functions is a pointer to the paired struct and this argument
/// can be used to track state in between calls (from an object oriented point
/// of view, this can be viewed as a "vtable" for a "class" -- that is the
/// corresponding struct above --; the first argument is in place of `this`).
///
/// Except where noted otherwise, all pointer arguments are owned by core
/// TensorFlow and are guaranteed to not be `nullptr`.
///
/// All path-like arguments are null terminated `char*` strings. Plugins can
/// assume that before any function using path arguments is invoked, the path is
/// made canonical by calling the function provided by `translate_name` or a
/// default implementation of that (supplied by core TensorFlow).
///
/// The only time the pointer to the `TF_*` structures from section 1 is not
/// marked `const` in these functions is when these function are either
/// allocating or deallocating the plugin specific data. That is, in the 4
/// `cleanup` functions (one for each data structure), the `init` function for
/// `TF_Filesystem` and the `new_*` methods of `TF_FilesystemOps` to initialize
/// the 3 types of files. In all other cases, there is no need to modify the
/// address of the opaque data pointer, hence the wrapper pointer is marked
/// `const`.
///
/// For consistency, the arguments on all these functions follow the same
/// pattern: first we have the opaque pointer argument ("this" above), then the
/// input arguments, then the in-out arguments (if any) and we finish the
/// argument list with the out arguments. We only use the return type for an out
/// parameter if that is a plain C type, as this ensures ABI compatibility
/// (returning structures has issues in case compiler options affect
/// optimizations such as RVO). If a status needs to be returned from these
/// methods, the last argument is always a `TF_Status *` (or an array of such
/// pointers) owned by core TensorFlow and guaranteed to not be `nullptr`.
///
/// To ensure ABI and API compatibility, we have out-of-bounds data that is used
/// by both core TensorFlow and the plugin at load time. We don't include this
/// data in the structures here to prevent cases when padding/packing enabled by
/// different compiler options breaks compatibility. For more details about how
/// this is used, please consult next sections. Here we just wrap these tables
/// in lint warnings so that changes here cause changes to the versioning data
/// as well. Here is a short summary of what changes are allowed:
/// * adding a new method at the end of a table is allowed at any time;
/// * any other change to these tables is only allowed on a major TensorFlow
/// version change (e.g., from 2.x to 3.0). This is provided as an escape
/// hatch to allow cleaning up these tables. Since any of these changes
/// break ABI compatibility and cause all plugins to be recompiled, these
/// type of changes should be extremely rare.
///
/// Next section will detail this as well as some corner cases that are out of
/// scope for now.
// LINT.IfChange
typedef struct TF_RandomAccessFileOps {
/// Releases resources associated with `*file`.
///
/// Requires that `*file` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_RandomAccessFile* file);
/// Reads up to `n` bytes from `*file` starting at `offset`.
///
/// The output is in `buffer`, core TensorFlow owns the buffer and guarantees
/// that at least `n` bytes are available.
///
/// Returns number of bytes read or -1 in case of error. Because of this
/// constraint and the fact that `ssize_t` is not defined in `stdint.h`/C++
/// standard, the return type is `int64_t`.
///
/// This is thread safe.
///
/// Note: the `buffer` argument is NOT a null terminated string!
///
/// Plugins:
/// * Must set `status` to `TF_OK` if exactly `n` bytes have been read.
/// * Must set `status` to `TF_OUT_OF_RANGE` if fewer than `n` bytes have
/// been read due to EOF.
/// * Must return -1 for any other error and must set `status` to any
/// other value to provide more information about the error.
int64_t (*read)(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status);
} TF_RandomAccessFileOps;
// LINT.ThenChange(:random_access_file_ops_version)
// LINT.IfChange
typedef struct TF_WritableFileOps {
/// Releases resources associated with `*file`.
///
/// Requires that `*file` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_WritableFile* file);
/// Appends `buffer` of size `n` to `*file`.
///
/// Core TensorFlow owns `buffer` and guarantees at least `n` bytes of storage
/// that can be used to write data.
///
/// Note: the `buffer` argument is NOT a null terminated string!
///
/// Plugins:
/// * Must set `status` to `TF_OK` if exactly `n` bytes have been written.
/// * Must set `status` to `TF_RESOURCE_EXHAUSTED` if fewer than `n` bytes
/// have been written, potentially due to quota/disk space.
/// * Might use any other error value for `status` to signal other errors.
void (*append)(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status);
/// Returns the current write position in `*file`.
///
/// Plugins should ensure that the implementation is idempotent, 2 identical
/// calls result in the same answer.
///
/// Plugins:
/// * Must set `status` to `TF_OK` and return current position if no error.
/// * Must set `status` to any other value and return -1 in case of error.
int64_t (*tell)(const TF_WritableFile* file, TF_Status* status);
/// Flushes `*file` and syncs contents to filesystem.
///
/// This call might not block, and when it returns the contents might not have
/// been fully persisted.
///
/// DEFAULT IMPLEMENTATION: No op.
void (*flush)(const TF_WritableFile* file, TF_Status* status);
/// Syncs contents of `*file` with the filesystem.
///
/// This call should block until filesystem confirms that all buffers have
/// been flushed and persisted.
///
/// DEFAULT IMPLEMENTATION: No op.
void (*sync)(const TF_WritableFile* file, TF_Status* status);
/// Closes `*file`.
///
/// Flushes all buffers and deallocates all resources.
///
/// Calling `close` must not result in calling `cleanup`.
///
/// Core TensorFlow will never call `close` twice.
void (*close)(const TF_WritableFile* file, TF_Status* status);
} TF_WritableFileOps;
// LINT.ThenChange(:writable_file_ops_version)
// LINT.IfChange
typedef struct TF_ReadOnlyMemoryRegionOps {
/// Releases resources associated with `*region`.
///
/// Requires that `*region` is not used in any concurrent or subsequent
/// operations.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_ReadOnlyMemoryRegion* region);
/// Returns a pointer to the memory region.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
const void* (*data)(const TF_ReadOnlyMemoryRegion* region);
/// Returns the length of the memory region in bytes.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
uint64_t (*length)(const TF_ReadOnlyMemoryRegion* region);
} TF_ReadOnlyMemoryRegionOps;
// LINT.ThenChange(:read_only_memory_region_ops_version)
// LINT.IfChange
typedef struct TF_FilesystemOps {
/// Acquires all resources used by the filesystem.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*init)(TF_Filesystem* filesystem, TF_Status* status);
/// Releases all resources used by the filesystem
///
/// NOTE: TensorFlow does not unload DSOs. Thus, the only way a filesystem
/// won't be registered anymore is if this function gets called by core
/// TensorFlow and the `TF_Filesystem*` object is destroyed. However, due to
/// registration being done in a static instance of `Env`, the destructor of
/// `FileSystem` is never called (see
/// https://github.com/tensorflow/tensorflow/issues/27535). In turn, this
/// function will never be called. There are plans to refactor registration
/// and fix this.
///
/// TODO(b/139060984): After all filesystems are converted, revisit note.
///
/// This operation must be provided. See "REQUIRED OPERATIONS" above.
void (*cleanup)(TF_Filesystem* filesystem);
/// Creates a new random access read-only file from given `path`.
///
/// After this call `file` may be concurrently accessed by multiple threads.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to an
/// existing file or one of the parent entries in `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid (e.g., malformed, or has a parent entry
/// which is a file).
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_RandomAccessFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_random_access_file)(const TF_Filesystem* filesystem,
const char* path, TF_RandomAccessFile* file,
TF_Status* status);
/// Creates an object to write to a file with the specified `path`.
///
/// If the file already exists, it is deleted and recreated. The `file` object
/// must only be accessed by one thread at a time.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_WritableFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_writable_file)(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
/// Creates an object to append to a file with the specified `path`.
///
/// If the file doesn't exists, it is first created with empty contents.
/// The `file` object must only be accessed by one thread at a time.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `file` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_WritableFileOps` table. See "REQUIRED OPERATIONS" above.
void (*new_appendable_file)(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
/// Creates a read-only region of memory from contents of `path`.
///
/// After this call `region` may be concurrently accessed by multiple threads.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `region` was updated.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to an
/// existing file or one of the parent entries in `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Must set `status` to `TF_INVALID_ARGUMENT` if `path` points to an
/// empty file.
/// * Might use any other error value for `status` to signal other errors.
///
/// REQUIREMENTS: If plugins implement this, they must also provide a filled
/// `TF_ReadOnlyMemoryRegionOps` table. See "REQUIRED OPERATIONS" above.
void (*new_read_only_memory_region_from_file)(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status);
/// Creates the directory specified by `path`, assuming parent exists.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was created.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Must set `status` to `TF_ALREADY_EXISTS` if `path` already exists.
/// * Might use any other error value for `status` to signal other errors.
void (*create_dir)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Creates the directory specified by `path` and all needed ancestors.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was created.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid or
/// if it exists but is not a directory.
/// * Might use any other error value for `status` to signal other errors.
///
/// NOTE: The requirements specify that `TF_ALREADY_EXISTS` is not returned if
/// directory exists. Similarly, `TF_NOT_FOUND` is not be returned, as the
/// missing directory entry and all its descendants will be created by the
/// plugin.
///
/// DEFAULT IMPLEMENTATION: Creates directories one by one. Needs
/// `path_exists`, `is_directory`, and `create_dir`.
void (*recursively_create_dir)(const TF_Filesystem* filesystem,
const char* path, TF_Status* status);
/// Deletes the file specified by `path`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if file was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` points to a
/// directory or if it is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*delete_file)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Deletes the empty directory specified by `path`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` does not point
/// to a directory, if `path` is invalid, or if directory is not empty.
/// * Might use any other error value for `status` to signal other errors.
void (*delete_dir)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Deletes the directory specified by `path` and all its contents.
///
/// This is accomplished by traversing directory tree rooted at `path` and
/// deleting entries as they are encountered, from leaves to root. Each plugin
/// is free to choose a different approach which obtains similar results.
///
/// On successful deletion, `status` must be `TF_OK` and `*undeleted_files`
/// and `*undeleted_dirs` must be 0. On unsuccessful deletion, `status` must
/// be set to the reason why one entry couldn't be removed and the proper
/// count must be updated. If the deletion is unsuccessful because the
/// traversal couldn't start, `*undeleted_files` must be set to 0 and
/// `*undeleted_dirs` must be set to 1.
///
/// TODO(b/139060984): After all filesystems are converted, consider
/// invariant about `*undeleted_files` and `*undeleted_dirs`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if directory was deleted.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Does a BFS traversal of tree rooted at `path`,
/// deleting entries as needed. Needs `path_exists`, `get_children`,
/// `is_directory`, `delete_file`, and `delete_dir`.
void (*delete_recursively)(const TF_Filesystem* filesystem, const char* path,
uint64_t* undeleted_files,
uint64_t* undeleted_dirs, TF_Status* status);
/// Renames the file given by `src` to that in `dst`.
///
/// Replaces `dst` if it exists. In case of error, both `src` and `dst` keep
/// the same state as before the call.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if rename was completed.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// either `src` or `dst` doesn't exist or if the specified `src` path
/// doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if either `src` or
/// `dst` is a directory or if either of them is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Copies file and deletes original. Needs
/// `copy_file`. and `delete_file`.
void (*rename_file)(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
/// Copies the file given by `src` to that in `dst`.
///
/// Similar to `rename_file`, but both `src` and `dst` exist after this call
/// with the same contents. In case of error, both `src` and `dst` keep the
/// same state as before the call.
///
/// If `dst` is a directory, creates a file with the same name as the source
/// inside the target directory.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if rename was completed.
/// * Must set `status` to `TF_NOT_FOUND` if one of the parents entries in
/// either `src` or `dst` doesn't exist or if the specified `src` path
/// doesn't exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if either `src` or
/// `dst` is a directory or if either of them is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Reads from `src` and writes to `dst`. Needs
/// `new_random_access_file` and `new_writable_file`.
void (*copy_file)(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
/// Checks if `path` exists.
///
/// Note that this doesn't differentiate between files and directories.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*path_exists)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Checks if all values in `paths` exist in the filesystem.
///
/// Returns `true` if and only if calling `path_exists` on each entry in
/// `paths` would set `status` to `TF_OK`.
///
/// Caller guarantees that:
/// * `paths` has exactly `num_files` entries.
/// * `statuses` is either null or an array of `num_files` non-null elements
/// of type `TF_Status*`.
///
/// If `statuses` is not null, plugins must fill each element with detailed
/// status for each file, as if calling `path_exists` on each one. Core
/// TensorFlow initializes the `statuses` array and plugins must use
/// `TF_SetStatus` to set each element instead of directly assigning.
///
/// DEFAULT IMPLEMENTATION: Checks existence of every file. Needs
/// `path_exists`.
bool (*paths_exist)(const TF_Filesystem* filesystem, char** paths,
int num_files, TF_Status** statuses);
/// Obtains statistics for the given `path`.
///
/// Updates `stats` only if `status` is set to `TF_OK`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
void (*stat)(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status);
/// Checks whether the given `path` is a directory or not.
///
/// If `status` is not `TF_OK`, returns `false`, otherwise returns the same
/// as the `is_directory` member of a `TF_FileStatistics` that would be used
/// on the equivalent call of `stat`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Gets statistics about `path`. Needs `stat`.
bool (*is_directory)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Returns the size of the file given by `path`.
///
/// If `status` is not `TF_OK`, return value is undefined. Otherwise, returns
/// the same as `length` member of a `TF_FileStatistics` that would be used on
/// the equivalent call of `stat`.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `path` exists.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if `path` is invalid or
/// points to a directory.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Gets statistics about `path`. Needs `stat`.
int64_t (*get_file_size)(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
/// Translates `uri` to a filename for the filesystem
///
/// A filesystem is registered for a specific scheme and all of the methods
/// should work with URIs. Hence, each filesystem needs to be able to
/// translate from an URI to a path on the filesystem. For example, this
/// function could translate `fs:///path/to/a/file` into `/path/to/a/file`, if
/// implemented by a filesystem registered to handle the `fs://` scheme.
///
/// A new `char*` buffer must be allocated by this method. Core TensorFlow
/// manages the lifetime of the buffer after the call. Thus, all callers of
/// this method must take ownership of the returned pointer.
///
/// The implementation should clean up paths, including but not limited to,
/// removing duplicate `/`s, and resolving `..` and `.`.
///
/// Plugins must not return `nullptr`. Returning empty strings is allowed.
///
/// The allocation and freeing of memory must happen via the functions sent to
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
/// structure in Section 4).
///
/// This function will be called by core TensorFlow to clean up all path
/// arguments for all other methods in the filesystem API.
///
/// DEFAULT IMPLEMENTATION: Uses `io::CleanPath` and `io::ParseURI`.
char* (*translate_name)(const TF_Filesystem* filesystem, const char* uri);
/// Finds all entries in the directory given by `path`.
///
/// The returned entries are paths relative to `path`.
///
/// Plugins must allocate `entries` to hold all names that need to be returned
/// and return the size of `entries`. Caller takes ownership of `entries`
/// after the call.
///
/// In case of error, plugins must set `status` to a value different than
/// `TF_OK`, free memory allocated for `entries` and return -1.
///
/// The allocation and freeing of memory must happen via the functions sent to
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
/// structure in Section 4).
///
/// Plugins:
/// * Must set `status` to `TF_OK` if all children were returned.
/// * Must set `status` to `TF_NOT_FOUND` if `path` doesn't point to a
/// filesystem entry or if one of the parents entries in `path` doesn't
/// exist.
/// * Must set `status` to `TF_FAILED_PRECONDITION` if one of the parent
/// entries in `path` is not a directory, or if `path` is a file.
/// * Might use any other error value for `status` to signal other errors.
int (*get_children)(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status);
/// Finds all entries matching the regular expression given by `glob`.
///
/// Pattern must match the entire entry name, not just a substring.
///
/// pattern: { term }
/// term:
/// '*': matches any sequence of non-'/' characters
/// '?': matches a single non-'/' character
/// '[' [ '^' ] { match-list } ']':
/// matches any single character (not) on the list
/// c: matches character c (c != '*', '?', '\\', '[')
/// '\\' c: matches character c
/// character-range:
/// c: matches character c (c != '\\', '-', ']')
/// '\\' c: matches character c
/// lo '-' hi: matches character c for lo <= c <= hi
///
/// Implementations must allocate `entries` to hold all names that need to be
/// returned and return the size of `entries`. Caller takes ownership of
/// `entries` after the call.
///
/// In case of error, the implementations must set `status` to a value
/// different than `TF_OK`, free any memory that might have been allocated for
/// `entries` and return -1.
///
/// The allocation and freeing of memory must happen via the functions sent to
/// core TensorFlow upon registration (see the `TF_FilesystemPluginInfo`
/// structure in Section 4).
///
/// Plugins:
/// * Must set `status` to `TF_OK` if all matches were returned.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: Scans the directory tree (in parallel if possible)
/// and fills `*entries`. Needs `get_children` and `is_directory`.
int (*get_matching_paths)(const TF_Filesystem* filesystem, const char* glob,
char*** entries, TF_Status* status);
/// Flushes any filesystem cache currently in memory
///
/// DEFAULT IMPLEMENTATION: No op.
void (*flush_caches)(const TF_Filesystem* filesystem);
/// Returns pointer to an array of available configuration options and their
/// current/default values in `options` and number of options in array in
/// `num_options`. Ownership of the array is transferred to caller and the
/// caller is responsible of freeing the buffers using respective file systems
/// allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `options` and `num_options` set.
/// If there is no configurable option, `num_options` should be 0.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return 0 options and `TF_OK`.
void (*get_filesystem_configuration)(const TF_Filesystem* filesystem,
TF_Filesystem_Option** options,
int* num_options, TF_Status* status);
/// Updates filesystem configuration with options passed in `options`. It can
/// contain full set of options supported by the filesystem or just a subset
/// of them. Ownership of options and buffers therein belongs to the caller
/// and any buffers need to be allocated through filesystem allocation API.
/// Filesystems may choose to ignore configuration errors but should at least
/// display a warning or error message to warn the users.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if options are updated.
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*set_filesystem_configuration)(const TF_Filesystem* filesystem,
const TF_Filesystem_Option* options,
int num_options, TF_Status* status);
/// Returns the value of the filesystem option given in `key` in `option`.
/// Valid values of the `key` are returned by
/// `get_file_system_configuration_keys` call. Ownership of the
/// `option` is transferred to caller. Buffers therein should be allocated and
/// freed by the relevant filesystems allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `option` is set
/// * Must set `status` to `TF_NOT_FOUND` if the key is invalid
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*get_filesystem_configuration_option)(const TF_Filesystem* filesystem,
const char* key,
TF_Filesystem_Option** option,
TF_Status* status);
/// Sets the value of the filesystem option given in `key` to value in
/// `option`. Valid values of the `key` are returned by
/// `get_file_system_configuration_keys` call. Ownership of the `option` and
/// the `key` belogs to the caller. Buffers therein should be allocated and
/// freed by the filesystems allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` if `option` is set/updated
/// * Must set `status` to `TF_NOT_FOUND` if the key is invalid
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_NOT_FOUND`.
void (*set_filesystem_configuration_option)(
const TF_Filesystem* filesystem, const TF_Filesystem_Option* option,
TF_Status* status);
/// Returns a list of valid configuration keys in `keys` array and number of
/// keys in `num_keys`. Ownership of the buffers in `keys` are transferred to
/// caller and needs to be freed using relevant filesystem allocation API.
///
/// Plugins:
/// * Must set `status` to `TF_OK` on success. If there are no configurable
/// keys, `num_keys` should be set to 0
/// * Might use any other error value for `status` to signal other errors.
///
/// DEFAULT IMPLEMENTATION: return `TF_OK` and `num_keys`=0.
void (*get_filesystem_configuration_keys)(const TF_Filesystem* filesystem,
char** keys, int* num_keys,
TF_Status* status);
} TF_FilesystemOps;
// LINT.ThenChange(:filesystem_ops_version)
/// SECTION 3. ABI and API compatibility
/// ----------------------------------------------------------------------------
///
/// In this section we define constants and macros to record versioning
/// information for each of the structures in section 2: ABI and API versions
/// and the number of functions in each of the function tables (which is
/// automatically determined, so ignored for the rest of this comment).
///
/// Since filesystem plugins are outside of TensorFlow's code tree, they are not
/// tied with TensorFlow releases and should have their own versioning metadata
/// in addition with the data discussed in this section. Each plugin author can
/// use a custom scheme, but it should only relate to changes in plugin code.
/// This section only touches metadata related to the versioning of this
/// interface that is shared by all possible plugins.
///
/// The API number increases whenever we break API compatibility while still
/// maintaining ABI compatibility. This happens only in the following cases:
/// 1. A new method is added _at the end_ of the function table.
/// 2. Preconditions or postconditions for one operation in these function
/// table change. Note that only core TensorFlow is able to impose these
/// invariants (i.e., guarantee the preconditions before calling the operation
/// and check the postconditions after the operation returns). If plugins need
/// additional invariants, they should be checked on the plugin side and the
/// `status` out variable should be updated accordingly (e.g., to include
/// plugin version information that relates to the condition change).
///
/// All other changes to the data structures (e.g., method removal, method
/// reordering, argument reordering, adding or removing arguments, changing the
/// type or the constness of a parameter, etc.) results in an ABI breakage.
/// Thus, we should not do any of these types of changes, except, potentially,
/// when we are releasing a new major version of TensorFlow. This is an escape
/// hatch, to be used rarely, preferably only to cleanup these structures.
/// Whenever we do these changes, the ABI number must be increased.
///
/// Next section will detail how this metadata is used at plugin registration to
/// only load compatible plugins and discard all others.
// LINT.IfChange(random_access_file_ops_version)
constexpr int TF_RANDOM_ACCESS_FILE_OPS_API = 0;
constexpr int TF_RANDOM_ACCESS_FILE_OPS_ABI = 0;
constexpr size_t TF_RANDOM_ACCESS_FILE_OPS_SIZE =
sizeof(TF_RandomAccessFileOps);
// LINT.ThenChange()
// LINT.IfChange(writable_file_ops_version)
constexpr int TF_WRITABLE_FILE_OPS_API = 0;
constexpr int TF_WRITABLE_FILE_OPS_ABI = 0;
constexpr size_t TF_WRITABLE_FILE_OPS_SIZE = sizeof(TF_WritableFileOps);
// LINT.ThenChange()
// LINT.IfChange(read_only_memory_region_ops_version)
constexpr int TF_READ_ONLY_MEMORY_REGION_OPS_API = 0;
constexpr int TF_READ_ONLY_MEMORY_REGION_OPS_ABI = 0;
constexpr size_t TF_READ_ONLY_MEMORY_REGION_OPS_SIZE =
sizeof(TF_ReadOnlyMemoryRegionOps);
// LINT.ThenChange()
// LINT.IfChange(filesystem_ops_version)
constexpr int TF_FILESYSTEM_OPS_API = 0;
constexpr int TF_FILESYSTEM_OPS_ABI = 0;
constexpr size_t TF_FILESYSTEM_OPS_SIZE = sizeof(TF_FilesystemOps);
// LINT.ThenChange()
/// SECTION 4. Plugin registration and initialization
/// ----------------------------------------------------------------------------
///
/// In this section we define the API used by core TensorFlow to initialize a
/// filesystem provided by a plugin. That is, we define the following:
/// * `TF_InitPlugin` function: must be present in the plugin shared object as
/// it will be called by core TensorFlow when the filesystem plugin is
/// loaded;
/// * `TF_FilesystemPluginOps` struct: used to transfer information between
/// plugins and core TensorFlow about the operations provided and metadata;
/// * `TF_FilesystemPluginInfo` struct: similar to the above structure, but
/// collects information about all the file schemes that the plugin provides
/// support for, as well as about the plugin's memory handling routines;
/// * `TF_SetFilesystemVersionMetadata` function: must be called by plugins in
/// their `TF_InitPlugin` to record the versioning information the plugins
/// are compiled against.
///
/// The `TF_InitPlugin` function is used by plugins to set up the data
/// structures that implement this interface, as presented in Section 2. In
/// order to not have plugin shared objects call back symbols defined in core
/// TensorFlow, `TF_InitPlugin` has a `TF_FilesystemPluginInfo` argument which
/// the plugin must fill (using the `TF_SetFilesystemVersionMetadata` for the
/// metadata and setting up all the supported operations and the URI schemes
/// that are supported).
/// This structure incorporates the operations defined in Section 2 and the
/// metadata defined in section 3, allowing plugins to define different ops
/// for different URI schemes.
///
/// Every URI scheme is of the form "fs" for URIs of form "fs:///path/to/file".
/// For local filesystems (i.e., when the URI is "/path/to/file"), the scheme
/// must be "". The scheme must never be `nullptr`.
///
/// Every plugin fills this in `TF_InitPlugin`, using the alocator passed as
/// argument to allocate memory. After `TF_InitPlugin` finishes, core
/// TensorFlow uses the information present in this to initialize filesystems
/// for the URI schemes that the plugin requests.
///
/// All pointers defined in this structure point to memory allocated by the DSO
/// using an allocator provided by core TensorFlow when calling `TF_InitPlugin`.
///
/// IMPORTANT: To maintain binary compatibility, the layout of this structure
/// must not change! In the unlikely case that a new type of file needs to be
/// supported, add the new ops and metadata at the end of the structure.
typedef struct TF_FilesystemPluginOps {
char* scheme;
int filesystem_ops_abi;
int filesystem_ops_api;
size_t filesystem_ops_size;
TF_FilesystemOps* filesystem_ops;
int random_access_file_ops_abi;
int random_access_file_ops_api;
size_t random_access_file_ops_size;
TF_RandomAccessFileOps* random_access_file_ops;
int writable_file_ops_abi;
int writable_file_ops_api;
size_t writable_file_ops_size;
TF_WritableFileOps* writable_file_ops;
int read_only_memory_region_ops_abi;
int read_only_memory_region_ops_api;
size_t read_only_memory_region_ops_size;
TF_ReadOnlyMemoryRegionOps* read_only_memory_region_ops;
} TF_FilesystemPluginOps;
/// This structure gathers together all the operations provided by the plugin.
///
/// Plugins must provide exactly `num_schemes` elements in the `ops` array.
///
/// Since memory that is allocated by the DSO gets transferred to core
/// TensorFlow, we need to provide a way for the allocation and deallocation to
/// match. This is why this structure also defines `plugin_memory_allocate` and
/// `plugin_memory_free` members.
///
/// All memory allocated by the plugin that will be owned by core TensorFlow
/// must be allocated using the allocator in this structure. Core TensorFlow
/// will use the deallocator to free this memory once it no longer needs it.
///
/// IMPORTANT: To maintain binary compatibility, the layout of this structure
/// must not change! In the unlikely case that new global operations must be
/// provided, add them at the end of the structure.
typedef struct TF_FilesystemPluginInfo {
size_t num_schemes;
TF_FilesystemPluginOps* ops;
void* (*plugin_memory_allocate)(size_t size);
void (*plugin_memory_free)(void* ptr);
} TF_FilesystemPluginInfo;
/// Convenience function for setting the versioning metadata.
///
/// The argument is guaranteed to not be `nullptr`.
///
/// We want this to be defined in the plugin's memory space and we guarantee
/// that core TensorFlow will never call this.
static inline void TF_SetFilesystemVersionMetadata(
TF_FilesystemPluginOps* ops) {
ops->filesystem_ops_abi = TF_FILESYSTEM_OPS_ABI;
ops->filesystem_ops_api = TF_FILESYSTEM_OPS_API;
ops->filesystem_ops_size = TF_FILESYSTEM_OPS_SIZE;
ops->random_access_file_ops_abi = TF_RANDOM_ACCESS_FILE_OPS_ABI;
ops->random_access_file_ops_api = TF_RANDOM_ACCESS_FILE_OPS_API;
ops->random_access_file_ops_size = TF_RANDOM_ACCESS_FILE_OPS_SIZE;
ops->writable_file_ops_abi = TF_WRITABLE_FILE_OPS_ABI;
ops->writable_file_ops_api = TF_WRITABLE_FILE_OPS_API;
ops->writable_file_ops_size = TF_WRITABLE_FILE_OPS_SIZE;
ops->read_only_memory_region_ops_abi = TF_READ_ONLY_MEMORY_REGION_OPS_ABI;
ops->read_only_memory_region_ops_api = TF_READ_ONLY_MEMORY_REGION_OPS_API;
ops->read_only_memory_region_ops_size = TF_READ_ONLY_MEMORY_REGION_OPS_SIZE;
}
/// Initializes a TensorFlow plugin.
///
/// Must be implemented by the plugin DSO. It is called by TensorFlow runtime.
///
/// Filesystem plugins can be loaded on demand by users via
/// `Env::LoadLibrary` or during TensorFlow's startup if they are on certain
/// paths (although this has a security risk if two plugins register for the
/// same filesystem and the malicious one loads before the legimitate one -
/// but we consider this to be something that users should care about and
/// manage themselves). In both of these cases, core TensorFlow looks for
/// the `TF_InitPlugin` symbol and calls this function.
///
/// For every filesystem URI scheme that this plugin supports, the plugin must
/// add one `TF_FilesystemPluginInfo` entry in `plugin_info->ops` and call
/// `TF_SetFilesystemVersionMetadata` for that entry.
///
/// Plugins must also initialize `plugin_info->plugin_memory_allocate` and
/// `plugin_info->plugin_memory_free` to ensure memory allocated by plugin is
/// freed in a compatible way.
TF_CAPI_EXPORT extern void TF_InitPlugin(TF_FilesystemPluginInfo* plugin_info);
#ifdef __cplusplus
} // end extern "C"
#endif // __cplusplus
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_FILESYSTEM_INTERFACE_H_
tensorflow/c/experimental/filesystem/modular_filesystem.cc
+577
View File
@@ -0,0 +1,577 @@
/* Copyright 2019 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/c/experimental/filesystem/modular_filesystem.h"
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/file_statistics.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/file_system_helper.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
// TODO(b/139060984): After all filesystems are converted, all calls to
// methods from `FileSystem` will have to be replaced to calls to private
// methods here, as part of making this class a singleton and the only way to
// register/use filesystems.
namespace tensorflow {
using UniquePtrTo_TF_Status =
::std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)>;
Status ModularFileSystem::NewRandomAccessFile(
const std::string& fname, std::unique_ptr<RandomAccessFile>* result) {
if (ops_->new_random_access_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewRandomAccessFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_RandomAccessFile>();
std::string translated_name = TranslateName(fname);
ops_->new_random_access_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularRandomAccessFile>(
translated_name, std::move(file), random_access_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewWritableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) {
if (ops_->new_writable_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewWritableFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_WritableFile>();
std::string translated_name = TranslateName(fname);
ops_->new_writable_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularWritableFile>(
translated_name, std::move(file), writable_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewAppendableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) {
if (ops_->new_appendable_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support NewAppendableFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto file = std::make_unique<TF_WritableFile>();
std::string translated_name = TranslateName(fname);
ops_->new_appendable_file(filesystem_.get(), translated_name.c_str(),
file.get(), plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularWritableFile>(
translated_name, std::move(file), writable_file_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::NewReadOnlyMemoryRegionFromFile(
const std::string& fname, std::unique_ptr<ReadOnlyMemoryRegion>* result) {
if (ops_->new_read_only_memory_region_from_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname,
" does not support NewReadOnlyMemoryRegionFromFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
auto region = std::make_unique<TF_ReadOnlyMemoryRegion>();
std::string translated_name = TranslateName(fname);
ops_->new_read_only_memory_region_from_file(
filesystem_.get(), translated_name.c_str(), region.get(),
plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK)
*result = std::make_unique<ModularReadOnlyMemoryRegion>(
std::move(region), read_only_memory_region_ops_.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::FileExists(absl::string_view fname) {
if (ops_->path_exists == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support FileExists()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
const std::string translated_name = TranslateName(fname);
ops_->path_exists(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
bool ModularFileSystem::FilesExist(const std::vector<std::string>& files,
std::vector<Status>* status) {
if (ops_->paths_exist == nullptr)
return FileSystem::FilesExist(files, status);
std::vector<char*> translated_names;
translated_names.reserve(files.size());
for (int i = 0; i < files.size(); i++)
translated_names.push_back(strdup(TranslateName(files[i]).c_str()));
bool result;
if (status == nullptr) {
result = ops_->paths_exist(filesystem_.get(), translated_names.data(),
files.size(), nullptr);
} else {
std::vector<TF_Status*> plugin_status;
plugin_status.reserve(files.size());
for (int i = 0; i < files.size(); i++)
plugin_status.push_back(TF_NewStatus());
result = ops_->paths_exist(filesystem_.get(), translated_names.data(),
files.size(), plugin_status.data());
for (int i = 0; i < files.size(); i++) {
status->push_back(StatusFromTF_Status(plugin_status[i]));
TF_DeleteStatus(plugin_status[i]);
}
}
for (int i = 0; i < files.size(); i++) free(translated_names[i]);
return result;
}
Status ModularFileSystem::GetChildren(const std::string& dir,
std::vector<std::string>* result) {
if (ops_->get_children == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dir, " does not support GetChildren()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dir);
// Note that `children` is allocated by the plugin and freed by core
// TensorFlow, so we need to use `plugin_memory_free_` here.
char** children = nullptr;
const int num_children =
ops_->get_children(filesystem_.get(), translated_name.c_str(), &children,
plugin_status.get());
if (num_children >= 0) {
for (int i = 0; i < num_children; i++) {
result->push_back(std::string(children[i]));
plugin_memory_free_(children[i]);
}
plugin_memory_free_(children);
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::GetMatchingPaths(const std::string& pattern,
std::vector<std::string>* result) {
if (ops_->get_matching_paths == nullptr)
return internal::GetMatchingPaths(this, Env::Default(), pattern, result);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
// Note that `matches` is allocated by the plugin and freed by core
// TensorFlow, so we need to use `plugin_memory_free_` here.
char** matches = nullptr;
const int num_matches = ops_->get_matching_paths(
filesystem_.get(), pattern.c_str(), &matches, plugin_status.get());
if (num_matches >= 0) {
for (int i = 0; i < num_matches; i++) {
result->push_back(std::string(matches[i]));
plugin_memory_free_(matches[i]);
}
plugin_memory_free_(matches);
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteFile(const std::string& fname) {
if (ops_->delete_file == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support DeleteFile()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
ops_->delete_file(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteRecursively(const std::string& dirname,
int64_t* undeleted_files,
int64_t* undeleted_dirs) {
if (undeleted_files == nullptr || undeleted_dirs == nullptr)
return errors::FailedPrecondition(
"DeleteRecursively must not be called with `undeleted_files` or "
"`undeleted_dirs` set to NULL");
if (ops_->delete_recursively == nullptr)
return FileSystem::DeleteRecursively(dirname, undeleted_files,
undeleted_dirs);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
uint64_t plugin_undeleted_files, plugin_undeleted_dirs;
ops_->delete_recursively(filesystem_.get(), translated_name.c_str(),
&plugin_undeleted_files, &plugin_undeleted_dirs,
plugin_status.get());
*undeleted_files = plugin_undeleted_files;
*undeleted_dirs = plugin_undeleted_dirs;
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::DeleteDir(const std::string& dirname) {
if (ops_->delete_dir == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dirname, " does not support DeleteDir()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->delete_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::RecursivelyCreateDir(const std::string& dirname) {
if (ops_->recursively_create_dir == nullptr)
return FileSystem::RecursivelyCreateDir(dirname);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->recursively_create_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::CreateDir(const std::string& dirname) {
if (ops_->create_dir == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", dirname, " does not support CreateDir()"));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(dirname);
ops_->create_dir(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::Stat(const std::string& fname, FileStatistics* stat) {
if (ops_->stat == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Filesystem for ", fname, " does not support Stat()"));
if (stat == nullptr)
return errors::InvalidArgument("FileStatistics pointer must not be NULL");
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
TF_FileStatistics stats;
ops_->stat(filesystem_.get(), translated_name.c_str(), &stats,
plugin_status.get());
if (TF_GetCode(plugin_status.get()) == TF_OK) {
stat->length = stats.length;
stat->mtime_nsec = stats.mtime_nsec;
stat->is_directory = stats.is_directory;
}
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::IsDirectory(const std::string& name) {
if (ops_->is_directory == nullptr) return FileSystem::IsDirectory(name);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(name);
ops_->is_directory(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::GetFileSize(const std::string& fname,
uint64* file_size) {
if (ops_->get_file_size == nullptr) {
FileStatistics stat;
Status status = Stat(fname, &stat);
if (!status.ok()) return status;
if (stat.is_directory)
return errors::FailedPrecondition("Called GetFileSize on a directory");
*file_size = stat.length;
return status;
}
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_name = TranslateName(fname);
*file_size = ops_->get_file_size(filesystem_.get(), translated_name.c_str(),
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::RenameFile(const std::string& src,
const std::string& target) {
if (ops_->rename_file == nullptr) {
Status status = CopyFile(src, target);
if (status.ok()) status = DeleteFile(src);
return status;
}
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_src = TranslateName(src);
std::string translated_target = TranslateName(target);
ops_->rename_file(filesystem_.get(), translated_src.c_str(),
translated_target.c_str(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::CopyFile(const std::string& src,
const std::string& target) {
if (ops_->copy_file == nullptr) return FileSystem::CopyFile(src, target);
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
std::string translated_src = TranslateName(src);
std::string translated_target = TranslateName(target);
ops_->copy_file(filesystem_.get(), translated_src.c_str(),
translated_target.c_str(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
std::string ModularFileSystem::TranslateName(absl::string_view name) const {
if (ops_->translate_name == nullptr) return FileSystem::TranslateName(name);
char* p = ops_->translate_name(filesystem_.get(), std::string(name).c_str());
CHECK(p != nullptr) << "TranslateName(" << name << ") returned nullptr";
std::string ret(p);
// Since `p` is allocated by plugin, free it using plugin's method.
plugin_memory_free_(p);
return ret;
}
void ModularFileSystem::FlushCaches() {
if (ops_->flush_caches != nullptr) ops_->flush_caches(filesystem_.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<string>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Buffer;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].buffer_val.buf = const_cast<char*>(values[i].c_str());
option_values[i].buffer_val.buf_length = values[i].size();
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<int64_t>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Int;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].int_val = values[i];
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularFileSystem::SetOption(const std::string& name,
const std::vector<double>& values) {
if (ops_->set_filesystem_configuration == nullptr) {
return errors::Unimplemented(
"Filesystem does not support SetConfiguration()");
}
if (values.empty()) {
return errors::InvalidArgument(
"SetConfiguration() needs number of values > 0");
}
TF_Filesystem_Option option;
memset(&option, 0, sizeof(option));
option.name = const_cast<char*>(name.c_str());
TF_Filesystem_Option_Value option_value;
memset(&option_value, 0, sizeof(option_value));
option_value.type_tag = TF_Filesystem_Option_Type_Real;
option_value.num_values = values.size();
std::vector<TF_Filesystem_Option_Value_Union> option_values(values.size());
for (size_t i = 0; i < values.size(); i++) {
memset(&option_values[i], 0, sizeof(option_values[i]));
option_values[i].real_val = values[i];
}
option_value.values = &option_values[0];
option.value = &option_value;
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->set_filesystem_configuration(filesystem_.get(), &option, 1,
plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularRandomAccessFile::Read(uint64 offset, size_t n,
StringPiece* result, char* scratch) const {
if (ops_->read == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Read() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
int64_t read =
ops_->read(file_.get(), offset, n, scratch, plugin_status.get());
if (read > 0) *result = StringPiece(scratch, read);
return StatusFromTF_Status(plugin_status.get());
}
Status ModularRandomAccessFile::Name(StringPiece* result) const {
*result = filename_;
return OkStatus();
}
Status ModularWritableFile::Append(StringPiece data) {
if (ops_->append == nullptr)
return errors::Unimplemented(tensorflow::strings::StrCat(
"Append() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->append(file_.get(), data.data(), data.size(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Close() {
if (ops_->close == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Close() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->close(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Flush() {
if (ops_->flush == nullptr) return OkStatus();
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->flush(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Sync() {
if (ops_->sync == nullptr) return Flush();
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
ops_->sync(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status ModularWritableFile::Name(StringPiece* result) const {
*result = filename_;
return OkStatus();
}
Status ModularWritableFile::Tell(int64_t* position) {
if (ops_->tell == nullptr)
return errors::Unimplemented(
tensorflow::strings::StrCat("Tell() not implemented for ", filename_));
UniquePtrTo_TF_Status plugin_status(TF_NewStatus(), TF_DeleteStatus);
*position = ops_->tell(file_.get(), plugin_status.get());
return StatusFromTF_Status(plugin_status.get());
}
Status RegisterFilesystemPlugin(const std::string& dso_path) {
// Step 1: Load plugin
Env* env = Env::Default();
void* dso_handle;
TF_RETURN_IF_ERROR(env->LoadDynamicLibrary(dso_path.c_str(), &dso_handle));
// Step 2: Load symbol for `TF_InitPlugin`
void* dso_symbol;
TF_RETURN_WITH_CONTEXT_IF_ERROR(
env->GetSymbolFromLibrary(dso_handle, "TF_InitPlugin", &dso_symbol),
"Failed to load TF_InitPlugin symbol for DSO: ", dso_path);
// Step 3: Call `TF_InitPlugin`
TF_FilesystemPluginInfo info;
memset(&info, 0, sizeof(info));
auto TF_InitPlugin =
reinterpret_cast<int (*)(TF_FilesystemPluginInfo*)>(dso_symbol);
TF_InitPlugin(&info);
// Step 4: Do the actual registration
return filesystem_registration::RegisterFilesystemPluginImpl(&info);
}
} // namespace tensorflow
tensorflow/c/experimental/filesystem/modular_filesystem.h
+197
View File
@@ -0,0 +1,197 @@
/* Copyright 2019 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_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "xla/tsl/platform/file_system.h"
#include "tensorflow/core/platform/file_statistics.h"
#include "tensorflow/core/platform/file_system.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
/// This file builds classes needed to hold a filesystem implementation in the
/// modular world. Once all TensorFlow filesystems are converted to use the
/// plugin based approach, this file will replace the one in core/platform and
/// the names will lose the `Modular` part. Until that point, the `Modular*`
/// classes here are experimental and subject to breaking changes.
/// For documentation on these methods, consult `core/platform/filesystem.h`.
namespace tensorflow {
// TODO(b/143949615): After all filesystems are converted, this file will be
// moved to core/platform, and this class can become a singleton and replace the
// need for `Env::Default()`. At that time, we might decide to remove the need
// for `Env::Default()` altogether, but that's a different project, not in
// scope for now. I'm just mentioning this here as that transition will mean
// removal of the registration part from `Env` and adding it here instead: we
// will need tables to hold for each scheme the function tables that implement
// the needed functionality instead of the current `FileSystemRegistry` code in
// `core/platform/env.cc`.
class ModularFileSystem final : public FileSystem {
public:
ModularFileSystem(
std::unique_ptr<TF_Filesystem> filesystem,
std::unique_ptr<const TF_FilesystemOps> filesystem_ops,
std::unique_ptr<const TF_RandomAccessFileOps> random_access_file_ops,
std::unique_ptr<const TF_WritableFileOps> writable_file_ops,
std::unique_ptr<const TF_ReadOnlyMemoryRegionOps>
read_only_memory_region_ops,
std::function<void*(size_t)> plugin_memory_allocate,
std::function<void(void*)> plugin_memory_free)
: filesystem_(std::move(filesystem)),
ops_(std::move(filesystem_ops)),
random_access_file_ops_(std::move(random_access_file_ops)),
writable_file_ops_(std::move(writable_file_ops)),
read_only_memory_region_ops_(std::move(read_only_memory_region_ops)),
plugin_memory_allocate_(std::move(plugin_memory_allocate)),
plugin_memory_free_(std::move(plugin_memory_free)) {}
~ModularFileSystem() override { ops_->cleanup(filesystem_.get()); }
absl::Status NewRandomAccessFile(
const std::string& fname,
std::unique_ptr<RandomAccessFile>* result) override;
absl::Status NewWritableFile(const std::string& fname,
std::unique_ptr<WritableFile>* result) override;
absl::Status NewAppendableFile(
const std::string& fname, std::unique_ptr<WritableFile>* result) override;
absl::Status NewReadOnlyMemoryRegionFromFile(
const std::string& fname,
std::unique_ptr<ReadOnlyMemoryRegion>* result) override;
absl::Status FileExists(absl::string_view fname) override;
bool FilesExist(const std::vector<std::string>& files,
std::vector<absl::Status>* status) override;
absl::Status GetChildren(const std::string& dir,
std::vector<std::string>* result) override;
absl::Status GetMatchingPaths(const std::string& pattern,
std::vector<std::string>* results) override;
absl::Status DeleteFile(const std::string& fname) override;
absl::Status DeleteRecursively(const std::string& dirname,
int64_t* undeleted_files,
int64_t* undeleted_dirs) override;
absl::Status DeleteDir(const std::string& dirname) override;
absl::Status RecursivelyCreateDir(const std::string& dirname) override;
absl::Status CreateDir(const std::string& dirname) override;
absl::Status Stat(const std::string& fname, FileStatistics* stat) override;
absl::Status IsDirectory(const std::string& fname) override;
absl::Status GetFileSize(const std::string& fname,
uint64_t* file_size) override;
absl::Status RenameFile(const std::string& src,
const std::string& target) override;
absl::Status CopyFile(const std::string& src,
const std::string& target) override;
std::string TranslateName(absl::string_view name) const override;
void FlushCaches() override;
absl::Status SetOption(const std::string& name,
const std::vector<std::string>& values) override;
absl::Status SetOption(const std::string& name,
const std::vector<int64_t>& values) override;
absl::Status SetOption(const std::string& name,
const std::vector<double>& values) override;
private:
std::unique_ptr<TF_Filesystem> filesystem_;
std::unique_ptr<const TF_FilesystemOps> ops_;
std::unique_ptr<const TF_RandomAccessFileOps> random_access_file_ops_;
std::unique_ptr<const TF_WritableFileOps> writable_file_ops_;
std::unique_ptr<const TF_ReadOnlyMemoryRegionOps>
read_only_memory_region_ops_;
std::function<void*(size_t)> plugin_memory_allocate_;
std::function<void(void*)> plugin_memory_free_;
ModularFileSystem(const ModularFileSystem&) = delete;
void operator=(const ModularFileSystem&) = delete;
};
class ModularRandomAccessFile final : public RandomAccessFile {
public:
ModularRandomAccessFile(const std::string& filename,
std::unique_ptr<TF_RandomAccessFile> file,
const TF_RandomAccessFileOps* ops)
: filename_(filename), file_(std::move(file)), ops_(ops) {}
~ModularRandomAccessFile() override { ops_->cleanup(file_.get()); }
absl::Status Read(uint64_t offset, size_t n, absl::string_view* result,
char* scratch) const override;
absl::Status Name(absl::string_view* result) const override;
private:
std::string filename_;
std::unique_ptr<TF_RandomAccessFile> file_;
const TF_RandomAccessFileOps* ops_; // not owned
ModularRandomAccessFile(const ModularRandomAccessFile&) = delete;
void operator=(const ModularRandomAccessFile&) = delete;
};
class ModularWritableFile final : public WritableFile {
public:
ModularWritableFile(const std::string& filename,
std::unique_ptr<TF_WritableFile> file,
const TF_WritableFileOps* ops)
: filename_(filename), file_(std::move(file)), ops_(ops) {}
~ModularWritableFile() override { ops_->cleanup(file_.get()); }
absl::Status Append(absl::string_view data) override;
absl::Status Close() override;
absl::Status Flush() override;
absl::Status Sync() override;
absl::Status Name(absl::string_view* result) const override;
absl::Status Tell(int64_t* position) override;
private:
std::string filename_;
std::unique_ptr<TF_WritableFile> file_;
const TF_WritableFileOps* ops_; // not owned
ModularWritableFile(const ModularWritableFile&) = delete;
void operator=(const ModularWritableFile&) = delete;
};
class ModularReadOnlyMemoryRegion final : public ReadOnlyMemoryRegion {
public:
ModularReadOnlyMemoryRegion(std::unique_ptr<TF_ReadOnlyMemoryRegion> region,
const TF_ReadOnlyMemoryRegionOps* ops)
: region_(std::move(region)), ops_(ops) {}
~ModularReadOnlyMemoryRegion() override { ops_->cleanup(region_.get()); };
const void* data() override { return ops_->data(region_.get()); }
uint64_t length() override { return ops_->length(region_.get()); }
private:
std::unique_ptr<TF_ReadOnlyMemoryRegion> region_;
const TF_ReadOnlyMemoryRegionOps* ops_; // not owned
ModularReadOnlyMemoryRegion(const ModularReadOnlyMemoryRegion&) = delete;
void operator=(const ModularReadOnlyMemoryRegion&) = delete;
};
// Registers a filesystem plugin so that core TensorFlow can use it.
absl::Status RegisterFilesystemPlugin(const std::string& dso_path);
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/modular_filesystem_registration.cc
+338
View File
@@ -0,0 +1,338 @@
/* 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/c/experimental/filesystem/modular_filesystem_registration.h"
#include <algorithm>
#include <cstddef>
#include <memory>
#include <utility>
#include "absl/log/log.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/stringpiece.h"
namespace tensorflow {
// Checks that all schemes provided by a plugin are valid.
// TODO(b/139060984): More validation could be done here, based on supported
// charset, maximum length, etc. Punting it for later.
static Status ValidateScheme(const char* scheme) {
if (scheme == nullptr)
return errors::InvalidArgument(
"Attempted to register filesystem with `nullptr` URI scheme");
return OkStatus();
}
// Checks if the plugin and core ABI numbers match.
//
// If the numbers don't match, plugin cannot be loaded.
static Status CheckABI(int pluginABI, int coreABI, StringPiece where) {
if (pluginABI != coreABI)
return errors::FailedPrecondition(
strings::StrCat("Plugin ABI (", pluginABI, ") for ", where,
" operations doesn't match expected core ABI (",
coreABI, "). Plugin cannot be loaded."));
return OkStatus();
}
// Checks if the plugin and core ABI numbers match, for all operations.
//
// If the numbers don't match, plugin cannot be loaded.
//
// Uses the simpler `CheckABI(int, int, StringPiece)`.
static Status ValidateABI(const TF_FilesystemPluginOps* ops) {
TF_RETURN_IF_ERROR(
CheckABI(ops->filesystem_ops_abi, TF_FILESYSTEM_OPS_ABI, "filesystem"));
if (ops->random_access_file_ops != nullptr)
TF_RETURN_IF_ERROR(CheckABI(ops->random_access_file_ops_abi,
TF_RANDOM_ACCESS_FILE_OPS_ABI,
"random access file"));
if (ops->writable_file_ops != nullptr)
TF_RETURN_IF_ERROR(CheckABI(ops->writable_file_ops_abi,
TF_WRITABLE_FILE_OPS_ABI, "writable file"));
if (ops->read_only_memory_region_ops != nullptr)
TF_RETURN_IF_ERROR(CheckABI(ops->read_only_memory_region_ops_abi,
TF_READ_ONLY_MEMORY_REGION_OPS_ABI,
"read only memory region"));
return OkStatus();
}
// Checks if the plugin and core API numbers match, logging mismatches.
static void CheckAPI(int plugin_API, int core_API, StringPiece where) {
if (plugin_API != core_API) {
VLOG(0) << "Plugin API (" << plugin_API << ") for " << where
<< " operations doesn't match expected core API (" << core_API
<< "). Plugin will be loaded but functionality might be missing.";
}
}
// Checks if the plugin and core API numbers match, for all operations.
//
// Uses the simpler `CheckAPIHelper(int, int, StringPiece)`.
static void ValidateAPI(const TF_FilesystemPluginOps* ops) {
CheckAPI(ops->filesystem_ops_api, TF_FILESYSTEM_OPS_API, "filesystem");
if (ops->random_access_file_ops != nullptr)
CheckAPI(ops->random_access_file_ops_api, TF_RANDOM_ACCESS_FILE_OPS_API,
"random access file");
if (ops->writable_file_ops != nullptr)
CheckAPI(ops->writable_file_ops_api, TF_WRITABLE_FILE_OPS_API,
"writable file");
if (ops->read_only_memory_region_ops != nullptr)
CheckAPI(ops->read_only_memory_region_ops_api,
TF_READ_ONLY_MEMORY_REGION_OPS_API, "read only memory region");
}
// Validates the filesystem operations supplied by the plugin.
static Status ValidateHelper(const TF_FilesystemOps* ops) {
if (ops == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without operations");
if (ops->init == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `init` operation");
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation");
return OkStatus();
}
// Validates the random access file operations supplied by the plugin.
static Status ValidateHelper(const TF_RandomAccessFileOps* ops) {
if (ops == nullptr) {
// We allow filesystems where files can only be written to (from TF code)
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on random "
"access files");
return OkStatus();
}
// Validates the writable file operations supplied by the plugin.
static Status ValidateHelper(const TF_WritableFileOps* ops) {
if (ops == nullptr) {
// We allow read-only filesystems
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on writable "
"files");
return OkStatus();
}
// Validates the read only memory region operations given by the plugin.
static Status ValidateHelper(const TF_ReadOnlyMemoryRegionOps* ops) {
if (ops == nullptr) {
// read only memory region support is always optional
return OkStatus();
}
if (ops->cleanup == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `cleanup` operation on read "
"only memory regions");
if (ops->data == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `data` operation on read only "
"memory regions");
if (ops->length == nullptr)
return errors::FailedPrecondition(
"Trying to register filesystem without `length` operation on read only "
"memory regions");
return OkStatus();
}
// Validates the operations supplied by the plugin.
//
// Uses the 4 simpler `ValidateHelper(const TF_...*)` to validate each
// individual function table and then checks that the function table for a
// specific file type exists if the plugin offers support for creating that
// type of files.
static Status ValidateOperations(const TF_FilesystemPluginOps* ops) {
TF_RETURN_IF_ERROR(ValidateHelper(ops->filesystem_ops));
TF_RETURN_IF_ERROR(ValidateHelper(ops->random_access_file_ops));
TF_RETURN_IF_ERROR(ValidateHelper(ops->writable_file_ops));
TF_RETURN_IF_ERROR(ValidateHelper(ops->read_only_memory_region_ops));
if (ops->filesystem_ops->new_random_access_file != nullptr &&
ops->random_access_file_ops == nullptr)
return errors::FailedPrecondition(
"Filesystem allows creation of random access files but no "
"operations on them have been supplied.");
if ((ops->filesystem_ops->new_writable_file != nullptr ||
ops->filesystem_ops->new_appendable_file != nullptr) &&
ops->writable_file_ops == nullptr)
return errors::FailedPrecondition(
"Filesystem allows creation of writable files but no "
"operations on them have been supplied.");
if (ops->filesystem_ops->new_read_only_memory_region_from_file != nullptr &&
ops->read_only_memory_region_ops == nullptr)
return errors::FailedPrecondition(
"Filesystem allows creation of readonly memory regions but no "
"operations on them have been supplied.");
return OkStatus();
}
// Copies a function table from plugin memory space to core memory space.
//
// This has three benefits:
// * allows having newer plugins than the current core TensorFlow: the
// additional entries in the plugin's table are just discarded;
// * allows having older plugins than the current core TensorFlow (though
// we are still warning users): the entries that core TensorFlow expects
// but plugins didn't provide will be set to `nullptr` values and core
// TensorFlow will know to not call these on behalf of users;
// * increased security as plugins will not be able to alter function table
// after loading up. Thus, malicious plugins can't alter functionality to
// probe for gadgets inside core TensorFlow. We can even protect the area
// of memory where the copies reside to not allow any more writes to it
// after all copies are created.
template <typename T>
static std::unique_ptr<const T> CopyToCore(const T* plugin_ops,
size_t plugin_size) {
if (plugin_ops == nullptr) return nullptr;
size_t copy_size = std::min(plugin_size, sizeof(T));
auto core_ops = std::make_unique<T>();
memset(core_ops.get(), 0, sizeof(T));
memcpy(core_ops.get(), plugin_ops, copy_size);
return core_ops;
}
// Registers one filesystem from the plugin.
//
// Must be called only with `index` a valid index in `info->ops`.
static Status RegisterFileSystem(const TF_FilesystemPluginInfo* info,
int index) {
// Step 1: Copy all the function tables to core TensorFlow memory space
auto core_filesystem_ops = CopyToCore<TF_FilesystemOps>(
info->ops[index].filesystem_ops, info->ops[index].filesystem_ops_size);
auto core_random_access_file_ops = CopyToCore<TF_RandomAccessFileOps>(
info->ops[index].random_access_file_ops,
info->ops[index].random_access_file_ops_size);
auto core_writable_file_ops =
CopyToCore<TF_WritableFileOps>(info->ops[index].writable_file_ops,
info->ops[index].writable_file_ops_size);
auto core_read_only_memory_region_ops =
CopyToCore<TF_ReadOnlyMemoryRegionOps>(
info->ops[index].read_only_memory_region_ops,
info->ops[index].read_only_memory_region_ops_size);
// Step 2: Initialize the opaque filesystem structure
auto filesystem = std::make_unique<TF_Filesystem>();
TF_Status* c_status = TF_NewStatus();
Status status = OkStatus();
core_filesystem_ops->init(filesystem.get(), c_status);
status = Status(c_status->status);
TF_DeleteStatus(c_status);
if (!status.ok()) return status;
// Step 3: Actual registration
return Env::Default()->RegisterFileSystem(
info->ops[index].scheme,
std::make_unique<tensorflow::ModularFileSystem>(
std::move(filesystem), std::move(core_filesystem_ops),
std::move(core_random_access_file_ops),
std::move(core_writable_file_ops),
std::move(core_read_only_memory_region_ops),
info->plugin_memory_allocate, info->plugin_memory_free));
}
// Registers filesystem at `index`, if plugin is providing valid information.
//
// Extracted to a separate function so that pointers inside `info` are freed
// by the caller regardless of whether validation/registration failed or not.
//
// Must be called only with `index` a valid index in `info->ops`.
static Status ValidateAndRegisterFilesystems(
const TF_FilesystemPluginInfo* info, int index) {
TF_RETURN_IF_ERROR(ValidateScheme(info->ops[index].scheme));
TF_RETURN_IF_ERROR(ValidateABI(&info->ops[index]));
ValidateAPI(&info->ops[index]); // we just warn on API number mismatch
TF_RETURN_IF_ERROR(ValidateOperations(&info->ops[index]));
TF_RETURN_IF_ERROR(RegisterFileSystem(info, index));
return OkStatus();
}
// Ensures that the plugin provides the required memory management operations.
static Status ValidatePluginMemoryRoutines(
const TF_FilesystemPluginInfo* info) {
if (info->plugin_memory_allocate == nullptr)
return errors::FailedPrecondition(
"Cannot load filesystem plugin which does not provide "
"`plugin_memory_allocate`");
if (info->plugin_memory_free == nullptr)
return errors::FailedPrecondition(
"Cannot load filesystem plugin which does not provide "
"`plugin_memory_free`");
return OkStatus();
}
namespace filesystem_registration {
Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info) {
TF_RETURN_IF_ERROR(ValidatePluginMemoryRoutines(info));
// Validate and register all filesystems
// Try to register as many filesystems as possible.
// Free memory once we no longer need it
Status status;
for (int i = 0; i < info->num_schemes; i++) {
status.Update(ValidateAndRegisterFilesystems(info, i));
info->plugin_memory_free(info->ops[i].scheme);
info->plugin_memory_free(info->ops[i].filesystem_ops);
info->plugin_memory_free(info->ops[i].random_access_file_ops);
info->plugin_memory_free(info->ops[i].writable_file_ops);
info->plugin_memory_free(info->ops[i].read_only_memory_region_ops);
}
info->plugin_memory_free(info->ops);
return status;
}
} // namespace filesystem_registration
} // namespace tensorflow
tensorflow/c/experimental/filesystem/modular_filesystem_registration.h
+34
View File
@@ -0,0 +1,34 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
#include "absl/status/status.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
namespace filesystem_registration {
// Implementation for filesystem registration
//
// Don't call this directly. Instead call `RegisterFilesystemPlugin`.
// Exposed only for static registration of local filesystems.
absl::Status RegisterFilesystemPluginImpl(const TF_FilesystemPluginInfo* info);
} // namespace filesystem_registration
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_MODULAR_FILESYSTEM_REGISTRATION_H_
tensorflow/c/experimental/filesystem/modular_filesystem_test.cc
+1803
View File
File diff suppressed because it is too large Load Diff
tensorflow/c/experimental/filesystem/plugins/gcs/BUILD
+139
View File
@@ -0,0 +1,139 @@
# Experimental gcs filesystem plugin.
load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# Filesystem implementation for GCS environments
tf_cc_shared_object(
name = "gcs_filesystem",
framework_so = [],
linkstatic = False,
per_os_targets = 1,
visibility = ["//visibility:public"],
deps = [":gcs_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "gcs_filesystem_impl",
srcs = ["gcs_filesystem.cc"],
hdrs = ["gcs_filesystem.h"],
copts = select({
"//conditions:default": [],
"//tensorflow:windows": get_win_copts(),
}),
deps = [
":expiring_lru_cache",
":gcs_helper",
":ram_file_block_cache",
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
"@com_github_googlecloudplatform_google_cloud_cpp//:storage_client",
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
],
)
cc_library(
name = "gcs_helper",
srcs = ["gcs_helper.cc"],
hdrs = ["gcs_helper.h"],
linkstatic = 1,
deps = [
"//tensorflow/c:env",
],
)
cc_library(
name = "cleanup",
hdrs = ["cleanup.h"],
)
cc_library(
name = "ram_file_block_cache",
srcs = ["ram_file_block_cache.cc"],
hdrs = ["ram_file_block_cache.h"],
deps = [
":cleanup",
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)
tf_cc_test(
name = "ram_file_block_cache_test",
size = "small",
srcs = ["ram_file_block_cache_test.cc"],
deps = [
":ram_file_block_cache",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform/cloud:now_seconds_env",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tf_cc_test(
name = "gcs_filesystem_test",
srcs = [
"gcs_filesystem_test.cc",
],
tags = [
"manual",
"notap",
],
deps = [
":gcs_filesystem_impl",
"//tensorflow/core/platform:path",
"//tensorflow/core/platform:stacktrace_handler",
"//tensorflow/core/platform:test",
"@com_github_googlecloudplatform_google_cloud_cpp//google/cloud:google_cloud_cpp_common",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "expiring_lru_cache",
hdrs = ["expiring_lru_cache.h"],
deps = [
"//tensorflow/c:env",
"//tensorflow/c:tf_status",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/synchronization",
],
)
tf_cc_test(
name = "expiring_lru_cache_test",
size = "small",
srcs = ["expiring_lru_cache_test.cc"],
deps = [
":expiring_lru_cache",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c:tf_status_internal",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform/cloud:now_seconds_env",
],
)
tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h
+109
View File
@@ -0,0 +1,109 @@
/* 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.
==============================================================================*/
// MakeCleanup(f) returns an RAII cleanup object that calls 'f' in its
// destructor. The easiest way to use MakeCleanup is with a lambda argument,
// capturing the return value in an 'auto' local variable. Most users will not
// need more sophisticated syntax than that.
//
// Example:
// void func() {
// FILE* fp = fopen("data.txt", "r");
// if (fp == nullptr) return;
// auto fp_cleaner = gtl::MakeCleanup([fp] { fclose(fp); });
// // No matter what, fclose(fp) will happen.
// DataObject d;
// while (ReadDataObject(fp, &d)) {
// if (d.IsBad()) {
// LOG(ERROR) << "Bad Data";
// return;
// }
// PushGoodData(d);
// }
// }
//
// You can use Cleanup<F> directly, instead of using MakeCleanup and auto,
// but there's rarely a reason to do that.
//
// You can call 'release()' on a Cleanup object to cancel the cleanup.
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
#include <type_traits>
#include <utility>
namespace tf_gcs_filesystem {
// A move-only RAII object that calls a stored cleanup functor when
// destroyed. Cleanup<F> is the return type of gtl::MakeCleanup(F).
template <typename F>
class Cleanup {
public:
Cleanup() : released_(true), f_() {}
template <typename G>
explicit Cleanup(G&& f) // NOLINT
: f_(std::forward<G>(f)) {} // NOLINT(build/c++11)
Cleanup(Cleanup&& src) // NOLINT
: released_(src.is_released()), f_(src.release()) {}
// Implicitly move-constructible from any compatible Cleanup<G>.
// The source will be released as if src.release() were called.
// A moved-from Cleanup can be safely destroyed or reassigned.
template <typename G>
Cleanup(Cleanup<G>&& src) // NOLINT
: released_(src.is_released()), f_(src.release()) {}
// Assignment to a Cleanup object behaves like destroying it
// and making a new one in its place, analogous to unique_ptr
// semantics.
Cleanup& operator=(Cleanup&& src) { // NOLINT
if (!released_) f_();
released_ = src.released_;
f_ = src.release();
return *this;
}
~Cleanup() {
if (!released_) f_();
}
// Releases the cleanup function instead of running it.
// Hint: use c.release()() to run early.
F release() {
released_ = true;
return std::move(f_);
}
bool is_released() const { return released_; }
private:
static_assert(!std::is_reference<F>::value, "F must not be a reference");
bool released_ = false;
F f_;
};
template <int&... ExplicitParameterBarrier, typename F,
typename DecayF = typename std::decay<F>::type>
Cleanup<DecayF> MakeCleanup(F&& f) {
return Cleanup<DecayF>(std::forward<F>(f));
}
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_CLEANUP_H_
tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h
+191
View File
@@ -0,0 +1,191 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
#include <functional>
#include <list>
#include <map>
#include <memory>
#include <string>
#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "tensorflow/c/env.h"
#include "tensorflow/c/tf_status.h"
namespace tf_gcs_filesystem {
/// \brief An LRU cache of string keys and arbitrary values, with configurable
/// max item age (in seconds) and max entries.
///
/// This class is thread safe.
template <typename T>
class ExpiringLRUCache {
public:
/// A `max_age` of 0 means that nothing is cached. A `max_entries` of 0 means
/// that there is no limit on the number of entries in the cache (however, if
/// `max_age` is also 0, the cache will not be populated).
ExpiringLRUCache(uint64_t max_age, size_t max_entries,
std::function<uint64_t()> timer_seconds = TF_NowSeconds)
: max_age_(max_age),
max_entries_(max_entries),
timer_seconds_(timer_seconds) {}
/// Insert `value` with key `key`. This will replace any previous entry with
/// the same key.
void Insert(const std::string& key, const T& value) {
if (max_age_ == 0) {
return;
}
absl::MutexLock lock(mu_);
InsertLocked(key, value);
}
// Delete the entry with key `key`. Return true if the entry was found for
// `key`, false if the entry was not found. In both cases, there is no entry
// with key `key` existed after the call.
bool Delete(const std::string& key) {
absl::MutexLock lock(mu_);
return DeleteLocked(key);
}
/// Look up the entry with key `key` and copy it to `value` if found. Returns
/// true if an entry was found for `key`, and its timestamp is not more than
/// max_age_ seconds in the past.
bool Lookup(const std::string& key, T* value) {
if (max_age_ == 0) {
return false;
}
absl::MutexLock lock(mu_);
return LookupLocked(key, value);
}
typedef std::function<void(const std::string&, T*, TF_Status*)> ComputeFunc;
/// Look up the entry with key `key` and copy it to `value` if found. If not
/// found, call `compute_func`. If `compute_func` set `status` to `TF_OK`,
/// store a copy of the output parameter in the cache, and another copy in
/// `value`.
void LookupOrCompute(const std::string& key, T* value,
const ComputeFunc& compute_func, TF_Status* status) {
if (max_age_ == 0) {
return compute_func(key, value, status);
}
// Note: we hold onto mu_ for the rest of this function. In practice, this
// is okay, as stat requests are typically fast, and concurrent requests are
// often for the same file. Future work can split this up into one lock per
// key if this proves to be a significant performance bottleneck.
absl::MutexLock lock(mu_);
if (LookupLocked(key, value)) {
return TF_SetStatus(status, TF_OK, "");
}
compute_func(key, value, status);
if (TF_GetCode(status) == TF_OK) {
InsertLocked(key, *value);
}
}
/// Clear the cache.
void Clear() {
absl::MutexLock lock(mu_);
cache_.clear();
lru_list_.clear();
}
/// Accessors for cache parameters.
uint64_t max_age() const { return max_age_; }
size_t max_entries() const { return max_entries_; }
private:
struct Entry {
/// The timestamp (seconds) at which the entry was added to the cache.
uint64_t timestamp;
/// The entry's value.
T value;
/// A list iterator pointing to the entry's position in the LRU list.
std::list<std::string>::iterator lru_iterator;
};
bool LookupLocked(const std::string& key, T* value)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
auto it = cache_.find(key);
if (it == cache_.end()) {
return false;
}
lru_list_.erase(it->second.lru_iterator);
if (timer_seconds_() - it->second.timestamp > max_age_) {
cache_.erase(it);
return false;
}
*value = it->second.value;
lru_list_.push_front(it->first);
it->second.lru_iterator = lru_list_.begin();
return true;
}
void InsertLocked(const std::string& key, const T& value)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
lru_list_.push_front(key);
Entry entry{timer_seconds_(), value, lru_list_.begin()};
auto insert = cache_.insert(std::make_pair(key, entry));
if (!insert.second) {
lru_list_.erase(insert.first->second.lru_iterator);
insert.first->second = entry;
} else if (max_entries_ > 0 && cache_.size() > max_entries_) {
cache_.erase(lru_list_.back());
lru_list_.pop_back();
}
}
bool DeleteLocked(const std::string& key) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
auto it = cache_.find(key);
if (it == cache_.end()) {
return false;
}
lru_list_.erase(it->second.lru_iterator);
cache_.erase(it);
return true;
}
/// The maximum age of entries in the cache, in seconds. A value of 0 means
/// that no entry is ever placed in the cache.
const uint64_t max_age_;
/// The maximum number of entries in the cache. A value of 0 means there is no
/// limit on entry count.
const size_t max_entries_;
/// The callback to read timestamps.
std::function<uint64_t()> timer_seconds_;
/// Guards access to the cache and the LRU list.
absl::Mutex mu_;
/// The cache (a map from string key to Entry).
std::map<std::string, Entry> cache_ ABSL_GUARDED_BY(mu_);
/// The LRU list of entries. The front of the list identifies the most
/// recently accessed entry.
std::list<std::string> lru_list_ ABSL_GUARDED_BY(mu_);
};
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_EXPIRING_LRU_CACHE_H_
tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache_test.cc
+216
View File
@@ -0,0 +1,216 @@
/* 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/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
#include <cstdint>
#include <memory>
#include <string>
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/cloud/now_seconds_env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
TEST(ExpiringLRUCacheTest, MaxAge) {
const std::string key = "a";
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
tf_gcs_filesystem::ExpiringLRUCache<int> cache(
1, 0, [&env]() { return env->NowSeconds(); });
env->SetNowSeconds(1);
// Verify that replacement of an existing element works, and updates the
// timestamp of the entry.
cache.Insert(key, 41);
env->SetNowSeconds(2);
cache.Insert(key, 42);
// 1 second after the most recent insertion, the entry is still valid.
env->SetNowSeconds(3);
int value = 0;
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 42);
// 2 seconds after the most recent insertion, the entry is no longer valid.
env->SetNowSeconds(4);
EXPECT_FALSE(cache.Lookup(key, &value));
// Re-insert the entry.
cache.Insert(key, 43);
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 43);
// The entry is valid 1 second after the insertion...
env->SetNowSeconds(5);
value = 0;
EXPECT_TRUE(cache.Lookup(key, &value));
EXPECT_EQ(value, 43);
// ...but is no longer valid 2 seconds after the insertion.
env->SetNowSeconds(6);
EXPECT_FALSE(cache.Lookup(key, &value));
}
TEST(ExpiringLRUCacheTest, MaxEntries) {
// max_age of 0 means nothing will be cached.
tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
cache1.Insert("a", 1);
int value = 0;
EXPECT_FALSE(cache1.Lookup("a", &value));
// Now set max_age = 1 and verify the LRU eviction logic.
tf_gcs_filesystem::ExpiringLRUCache<int> cache2(1, 4);
cache2.Insert("a", 1);
cache2.Insert("b", 2);
cache2.Insert("c", 3);
cache2.Insert("d", 4);
EXPECT_TRUE(cache2.Lookup("a", &value));
EXPECT_EQ(value, 1);
EXPECT_TRUE(cache2.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache2.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache2.Lookup("d", &value));
EXPECT_EQ(value, 4);
// Insertion of "e" causes "a" to be evicted, but the other entries are still
// there.
cache2.Insert("e", 5);
EXPECT_FALSE(cache2.Lookup("a", &value));
EXPECT_TRUE(cache2.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache2.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache2.Lookup("d", &value));
EXPECT_EQ(value, 4);
EXPECT_TRUE(cache2.Lookup("e", &value));
EXPECT_EQ(value, 5);
}
TEST(ExpiringLRUCacheTest, LookupOrCompute) {
// max_age of 0 means we should always compute.
uint64_t num_compute_calls = 0;
tf_gcs_filesystem::ExpiringLRUCache<int>::ComputeFunc compute_func =
[&num_compute_calls](const std::string& key, int* value,
TF_Status* status) {
*value = num_compute_calls;
num_compute_calls++;
return TF_SetStatus(status, TF_OK, "");
};
tf_gcs_filesystem::ExpiringLRUCache<int> cache1(0, 4);
int value = -1;
TF_Status status;
cache1.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// re-read the same value, expect another lookup
cache1.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 2);
// Define a new cache with max_age > 0 and verify correct behavior.
tf_gcs_filesystem::ExpiringLRUCache<int> cache2(2, 4);
num_compute_calls = 0;
value = -1;
// Read our first value
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// Re-read, exepct no additional function compute_func calls.
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 0);
EXPECT_EQ(num_compute_calls, 1);
// Read a sequence of additional values, eventually evicting "a".
cache2.LookupOrCompute("b", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 2);
cache2.LookupOrCompute("c", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 2);
EXPECT_EQ(num_compute_calls, 3);
cache2.LookupOrCompute("d", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 3);
EXPECT_EQ(num_compute_calls, 4);
cache2.LookupOrCompute("e", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 4);
EXPECT_EQ(num_compute_calls, 5);
// Verify the other values remain in the cache.
cache2.LookupOrCompute("b", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 1);
EXPECT_EQ(num_compute_calls, 5);
cache2.LookupOrCompute("c", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 2);
EXPECT_EQ(num_compute_calls, 5);
cache2.LookupOrCompute("d", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 3);
EXPECT_EQ(num_compute_calls, 5);
// Re-read "a", ensure it is re-computed.
cache2.LookupOrCompute("a", &value, compute_func, &status);
TF_EXPECT_OK(status.status);
EXPECT_EQ(value, 5);
EXPECT_EQ(num_compute_calls, 6);
}
TEST(ExpiringLRUCacheTest, Clear) {
tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
cache.Insert("a", 1);
cache.Insert("b", 2);
cache.Insert("c", 3);
cache.Insert("d", 4);
int value = 0;
EXPECT_TRUE(cache.Lookup("a", &value));
EXPECT_EQ(value, 1);
EXPECT_TRUE(cache.Lookup("b", &value));
EXPECT_EQ(value, 2);
EXPECT_TRUE(cache.Lookup("c", &value));
EXPECT_EQ(value, 3);
EXPECT_TRUE(cache.Lookup("d", &value));
EXPECT_EQ(value, 4);
cache.Clear();
EXPECT_FALSE(cache.Lookup("a", &value));
EXPECT_FALSE(cache.Lookup("b", &value));
EXPECT_FALSE(cache.Lookup("c", &value));
EXPECT_FALSE(cache.Lookup("d", &value));
}
TEST(ExpiringLRUCacheTest, Delete) {
// Insert an entry.
tf_gcs_filesystem::ExpiringLRUCache<int> cache(1, 4);
cache.Insert("a", 1);
int value = 0;
EXPECT_TRUE(cache.Lookup("a", &value));
EXPECT_EQ(value, 1);
// Delete the entry.
EXPECT_TRUE(cache.Delete("a"));
EXPECT_FALSE(cache.Lookup("a", &value));
// Try deleting the entry again.
EXPECT_FALSE(cache.Delete("a"));
EXPECT_FALSE(cache.Lookup("a", &value));
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.cc
+1240
View File
File diff suppressed because it is too large Load Diff
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h
+117
View File
@@ -0,0 +1,117 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include "absl/base/thread_annotations.h"
#include "absl/synchronization/mutex.h"
#include "google/cloud/storage/client.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/expiring_lru_cache.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include "tensorflow/c/tf_status.h"
void ParseGCSPath(const std::string& fname, bool object_empty_ok,
std::string* bucket, std::string* object, TF_Status* status);
namespace tf_random_access_file {
void Cleanup(TF_RandomAccessFile* file);
int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status);
} // namespace tf_random_access_file
namespace tf_writable_file {
void Cleanup(TF_WritableFile* file);
void Append(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status);
int64_t Tell(const TF_WritableFile* file, TF_Status* status);
void Flush(const TF_WritableFile* file, TF_Status* status);
void Sync(const TF_WritableFile* file, TF_Status* status);
void Close(const TF_WritableFile* file, TF_Status* status);
} // namespace tf_writable_file
namespace tf_read_only_memory_region {
void Cleanup(TF_ReadOnlyMemoryRegion* region);
const void* Data(const TF_ReadOnlyMemoryRegion* region);
uint64_t Length(const TF_ReadOnlyMemoryRegion* region);
} // namespace tf_read_only_memory_region
namespace tf_gcs_filesystem {
typedef struct GcsFileStat {
TF_FileStatistics base;
int64_t generation_number;
} GcsFileStat;
typedef struct GCSFile {
google::cloud::storage::Client gcs_client; // owned
bool compose;
absl::Mutex block_cache_lock;
std::shared_ptr<RamFileBlockCache> file_block_cache
ABSL_GUARDED_BY(block_cache_lock);
uint64_t block_size; // Reads smaller than block_size will trigger a read
// of block_size.
std::unique_ptr<ExpiringLRUCache<GcsFileStat>> stat_cache;
GCSFile(google::cloud::storage::Client&& gcs_client);
// This constructor is used for testing purpose only.
GCSFile(google::cloud::storage::Client&& gcs_client, bool compose,
uint64_t block_size, size_t max_bytes, uint64_t max_staleness,
uint64_t stat_cache_max_age, size_t stat_cache_max_entries);
} GCSFile;
// This function is used to initialize a filesystem without the need of setting
// manually environement variables.
void InitTest(TF_Filesystem* filesystem, bool compose, uint64_t block_size,
size_t max_bytes, uint64_t max_staleness,
uint64_t stat_cache_max_age, size_t stat_cache_max_entries,
TF_Status* status);
void Init(TF_Filesystem* filesystem, TF_Status* status);
void Cleanup(TF_Filesystem* filesystem);
void NewRandomAccessFile(const TF_Filesystem* filesystem, const char* path,
TF_RandomAccessFile* file, TF_Status* status);
void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
void NewAppendableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status);
void NewReadOnlyMemoryRegionFromFile(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status);
int64_t GetFileSize(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void PathExists(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void CreateDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
int GetChildren(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status);
void DeleteFile(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void Stat(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status);
void DeleteDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status);
void CopyFile(const TF_Filesystem* filesystem, const char* src, const char* dst,
TF_Status* status);
void RenameFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status);
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_filesystem_test.cc
+705
View File
@@ -0,0 +1,705 @@
/* 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/c/experimental/filesystem/plugins/gcs/gcs_filesystem.h"
#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <iterator>
#include <memory>
#include <random>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "absl/strings/match.h"
#include "absl/strings/string_view.h"
#include "third_party/cloud_cpp/google/cloud/common_options.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/stacktrace_handler.h"
#include "tensorflow/core/platform/test.h"
#define ASSERT_TF_OK(x) ASSERT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
#define EXPECT_TF_OK(x) EXPECT_EQ(TF_OK, TF_GetCode(x)) << TF_Message(x)
static const char* content = "abcdefghijklmnopqrstuvwxyz1234567890";
// We will work with content_view instead of content.
static const absl::string_view content_view = content;
namespace gcs = google::cloud::storage;
static std::string InitializeTmpDir() {
// This env should be something like `gs://bucket/path`
const char* test_dir = getenv("GCS_TEST_TMPDIR");
if (test_dir != nullptr) {
std::string bucket, object;
TF_Status* status = TF_NewStatus();
ParseGCSPath(test_dir, true, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK) {
TF_DeleteStatus(status);
return "";
}
TF_DeleteStatus(status);
// We add a random value into `test_dir` to ensures that two consecutive
// runs are unlikely to clash.
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<> distribution;
std::string rng_val = std::to_string(distribution(gen));
return tensorflow::io::JoinPath(std::string(test_dir), rng_val);
} else {
return "";
}
}
static std::string* GetTmpDir() {
static std::string tmp_dir = InitializeTmpDir();
if (tmp_dir == "")
return nullptr;
else
return &tmp_dir;
}
namespace tensorflow {
namespace {
// TODO(vnvo2409): Refactor `gcs_filesystem_test` to remove unnecessary tests
// after porting all tests from
// `//tensorflow/core/platform/cloud:gcs_file_system_test`.
class GCSFilesystemTest : public ::testing::Test {
public:
void SetUp() override {
root_dir_ = io::JoinPath(
*GetTmpDir(),
::testing::UnitTest::GetInstance()->current_test_info()->name());
status_ = TF_NewStatus();
filesystem_ = new TF_Filesystem;
filesystem_->plugin_filesystem = nullptr;
// Because different tests requires different setup for filesystem. We
// initialize filesystem in each testcase.
}
void TearDown() override {
TF_DeleteStatus(status_);
if (filesystem_->plugin_filesystem != nullptr)
tf_gcs_filesystem::Cleanup(filesystem_);
delete filesystem_;
}
std::string GetURIForPath(absl::string_view path) {
const std::string translated_name =
tensorflow::io::JoinPath(root_dir_, path);
return translated_name;
}
std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile* file)>
GetWriter() {
std::unique_ptr<TF_WritableFile, void (*)(TF_WritableFile * file)> writer(
new TF_WritableFile, [](TF_WritableFile* file) {
if (file != nullptr) {
if (file->plugin_file != nullptr) tf_writable_file::Cleanup(file);
delete file;
}
});
writer->plugin_file = nullptr;
return writer;
}
std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile* file)>
GetReader() {
std::unique_ptr<TF_RandomAccessFile, void (*)(TF_RandomAccessFile * file)>
reader(new TF_RandomAccessFile, [](TF_RandomAccessFile* file) {
if (file != nullptr) {
if (file->plugin_file != nullptr)
tf_random_access_file::Cleanup(file);
delete file;
}
});
reader->plugin_file = nullptr;
return reader;
}
void WriteString(const std::string& path, const std::string& content) {
auto writer = GetWriter();
tf_gcs_filesystem::NewWritableFile(filesystem_, path.c_str(), writer.get(),
status_);
if (TF_GetCode(status_) != TF_OK) return;
tf_writable_file::Append(writer.get(), content.c_str(), content.length(),
status_);
if (TF_GetCode(status_) != TF_OK) return;
tf_writable_file::Close(writer.get(), status_);
if (TF_GetCode(status_) != TF_OK) return;
}
std::string ReadAll(const std::string& path) {
auto reader = GetReader();
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(),
reader.get(), status_);
if (TF_GetCode(status_) != TF_OK) return "";
auto file_size =
tf_gcs_filesystem::GetFileSize(filesystem_, path.c_str(), status_);
if (TF_GetCode(status_) != TF_OK) return "";
std::string content;
content.resize(file_size);
auto read = tf_random_access_file::Read(reader.get(), 0, file_size,
&content[0], status_);
if (TF_GetCode(status_) != TF_OK) return "";
if (read >= 0) content.resize(read);
if (file_size != content.size())
TF_SetStatus(
status_, TF_DATA_LOSS,
std::string("expected " + std::to_string(file_size) + " got " +
std::to_string(content.size()) + " bytes")
.c_str());
return content;
}
protected:
TF_Filesystem* filesystem_;
TF_Status* status_;
private:
std::string root_dir_;
};
::testing::AssertionResult WriteToServer(const std::string& path, size_t offset,
size_t length, gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto writer = gcs_client->WriteObject(bucket, object);
writer.write(content + offset, length);
writer.Close();
if (writer.metadata()) {
return ::testing::AssertionSuccess();
} else {
return ::testing::AssertionFailure()
<< writer.metadata().status().message();
}
}
::testing::AssertionResult InsertObject(const std::string& path,
const std::string& content,
gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto metadata = gcs_client->InsertObject(bucket, object, content);
if (metadata)
return ::testing::AssertionSuccess();
else
return ::testing::AssertionFailure() << metadata.status().message();
}
::testing::AssertionResult CompareSubString(int64_t offset, size_t length,
absl::string_view result,
size_t read) {
// Result isn't a null-terminated string so we have to wrap it inside a
// `string_view`
if (length == read && content_view.substr(offset, length) ==
absl::string_view(result).substr(0, read))
return ::testing::AssertionSuccess();
else
return ::testing::AssertionFailure()
<< "Result: " << absl::string_view(result).substr(0, read)
<< " Read: " << read;
}
::testing::AssertionResult CompareWithServer(const std::string& path,
size_t offset, size_t length,
gcs::Client* gcs_client,
TF_Status* status) {
std::string bucket, object;
ParseGCSPath(path, false, &bucket, &object, status);
if (TF_GetCode(status) != TF_OK)
return ::testing::AssertionFailure() << TF_Message(status);
auto reader = gcs_client->ReadObject(bucket, object);
if (!reader) {
return ::testing::AssertionFailure() << reader.status().message();
} else {
std::string content{std::istreambuf_iterator<char>{reader}, {}};
return CompareSubString(offset, length, content, content.length());
}
}
TEST_F(GCSFilesystemTest, ParseGCSPath) {
std::string bucket, object;
ParseGCSPath("gs://bucket/path/to/object", false, &bucket, &object, status_);
ASSERT_TF_OK(status_);
ASSERT_EQ(bucket, "bucket");
ASSERT_EQ(object, "path/to/object");
ParseGCSPath("gs://bucket/", true, &bucket, &object, status_);
ASSERT_TF_OK(status_);
ASSERT_EQ(bucket, "bucket");
ParseGCSPath("bucket/path/to/object", false, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
// bucket name must end with "/"
ParseGCSPath("gs://bucket", true, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
ParseGCSPath("gs://bucket/", false, &bucket, &object, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT);
}
TEST_F(GCSFilesystemTest, RandomAccessFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string filepath = GetURIForPath("a_file");
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
char* result = new char[content_view.length()];
int64_t read = tf_random_access_file::Read(file, 0, 1, result, status_);
ASSERT_EQ(read, -1) << "Read: " << read;
ASSERT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(WriteToServer(filepath, 0, content_view.length(),
&gcs_file->gcs_client, status_));
read = tf_random_access_file::Read(file, 0, content_view.length(), result,
status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(CompareSubString(0, content_view.length(), result, read));
read = tf_random_access_file::Read(file, 0, 4, result, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(CompareSubString(0, 4, result, read));
read = tf_random_access_file::Read(file, content_view.length() - 2, 4, result,
status_);
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
ASSERT_TRUE(CompareSubString(content_view.length() - 2, 2, result, read));
delete[] result;
tf_random_access_file::Cleanup(file);
delete file;
}
TEST_F(GCSFilesystemTest, WritableFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string filepath = GetURIForPath("a_file");
TF_WritableFile* file = new TF_WritableFile;
tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
tf_writable_file::Append(file, content, 4, status_);
ASSERT_TF_OK(status_);
auto length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 4);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
tf_writable_file::Append(file, content + 4, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 8);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
tf_writable_file::Close(file, status_);
ASSERT_TF_OK(status_);
tf_writable_file::Cleanup(file);
// Testing for compose objects
gcs_file->compose = true;
filepath = GetURIForPath("b_file");
tf_gcs_filesystem::NewWritableFile(filesystem_, filepath.c_str(), file,
status_);
ASSERT_TF_OK(status_);
tf_writable_file::Append(file, content, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 4);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 4, &gcs_file->gcs_client, status_));
tf_writable_file::Append(file, content + 4, 4, status_);
ASSERT_TF_OK(status_);
length = tf_writable_file::Tell(file, status_);
ASSERT_EQ(length, 8);
ASSERT_TF_OK(status_);
tf_writable_file::Flush(file, status_);
ASSERT_TF_OK(status_);
ASSERT_TRUE(
CompareWithServer(filepath, 0, 8, &gcs_file->gcs_client, status_));
tf_writable_file::Close(file, status_);
ASSERT_TF_OK(status_);
tf_writable_file::Cleanup(file);
delete file;
}
TEST_F(GCSFilesystemTest, ReadOnlyMemoryRegion) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(WriteToServer(path, 0, 0, &gcs_file->gcs_client, status_));
TF_ReadOnlyMemoryRegion* region = new TF_ReadOnlyMemoryRegion;
tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
region, status_);
ASSERT_EQ(TF_GetCode(status_), TF_INVALID_ARGUMENT) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
ASSERT_TRUE(WriteToServer(path, 0, content_view.length(),
&gcs_file->gcs_client, status_));
tf_gcs_filesystem::NewReadOnlyMemoryRegionFromFile(filesystem_, path.c_str(),
region, status_);
ASSERT_TF_OK(status_);
auto length = tf_read_only_memory_region::Length(region);
ASSERT_EQ(length, content_view.length());
auto data =
static_cast<const char*>(tf_read_only_memory_region::Data(region));
ASSERT_TRUE(CompareSubString(0, content_view.length(), data, length));
tf_read_only_memory_region::Cleanup(region);
delete region;
}
TEST_F(GCSFilesystemTest, PathExists) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("PathExists");
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_EQ(TF_NOT_FOUND, TF_GetCode(status_)) << TF_Message(status_);
TF_SetStatus(status_, TF_OK, "");
WriteString(path, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_TF_OK(status_);
}
TEST_F(GCSFilesystemTest, GetChildren) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string base = GetURIForPath("GetChildren");
tf_gcs_filesystem::CreateDir(filesystem_, base.c_str(), status_);
EXPECT_TF_OK(status_);
const std::string file = io::JoinPath(base, "TestFile.csv");
WriteString(file, "test");
EXPECT_TF_OK(status_);
const std::string subdir = io::JoinPath(base, "SubDir");
tf_gcs_filesystem::CreateDir(filesystem_, subdir.c_str(), status_);
EXPECT_TF_OK(status_);
const std::string subfile = io::JoinPath(subdir, "TestSubFile.csv");
WriteString(subfile, "test");
EXPECT_TF_OK(status_);
char** entries;
auto num_entries = tf_gcs_filesystem::GetChildren(filesystem_, base.c_str(),
&entries, status_);
EXPECT_TF_OK(status_);
std::vector<std::string> childrens;
for (int i = 0; i < num_entries; ++i) {
childrens.push_back(entries[i]);
}
std::sort(childrens.begin(), childrens.end());
EXPECT_EQ(std::vector<std::string>({"SubDir/", "TestFile.csv"}), childrens);
}
TEST_F(GCSFilesystemTest, DeleteFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("DeleteFile");
WriteString(path, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::DeleteFile(filesystem_, path.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, path.c_str(), status_);
EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND);
}
TEST_F(GCSFilesystemTest, CreateDir) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string dir = GetURIForPath("CreateDir");
tf_gcs_filesystem::CreateDir(filesystem_, dir.c_str(), status_);
EXPECT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
EXPECT_TF_OK(status_);
EXPECT_TRUE(stat.is_directory);
}
TEST_F(GCSFilesystemTest, DeleteDir) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string dir = GetURIForPath("DeleteDir");
const std::string file = io::JoinPath(dir, "DeleteDirFile.csv");
WriteString(file, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
EXPECT_EQ(TF_GetCode(status_), TF_FAILED_PRECONDITION);
TF_SetStatus(status_, TF_OK, "");
tf_gcs_filesystem::DeleteFile(filesystem_, file.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::DeleteDir(filesystem_, dir.c_str(), status_);
EXPECT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, dir.c_str(), &stat, status_);
EXPECT_EQ(TF_GetCode(status_), TF_NOT_FOUND) << TF_Message(status_);
}
TEST_F(GCSFilesystemTest, StatFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string path = GetURIForPath("StatFile");
WriteString(path, "test");
ASSERT_TF_OK(status_);
TF_FileStatistics stat;
tf_gcs_filesystem::Stat(filesystem_, path.c_str(), &stat, status_);
EXPECT_TF_OK(status_);
EXPECT_EQ(4, stat.length);
EXPECT_FALSE(stat.is_directory);
}
TEST_F(GCSFilesystemTest, RenameFile) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string src = GetURIForPath("RenameFileSrc");
const std::string dst = GetURIForPath("RenameFileDst");
WriteString(src, "test");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
EXPECT_TF_OK(status_);
auto result = ReadAll(dst);
EXPECT_TF_OK(status_);
EXPECT_EQ("test", result);
}
TEST_F(GCSFilesystemTest, RenameFileOverwrite) {
tf_gcs_filesystem::Init(filesystem_, status_);
ASSERT_TF_OK(status_);
const std::string src = GetURIForPath("RenameFileOverwriteSrc");
const std::string dst = GetURIForPath("RenameFileOverwriteDst");
WriteString(src, "test_old");
ASSERT_TF_OK(status_);
WriteString(dst, "test_new");
ASSERT_TF_OK(status_);
tf_gcs_filesystem::PathExists(filesystem_, dst.c_str(), status_);
EXPECT_TF_OK(status_);
tf_gcs_filesystem::RenameFile(filesystem_, src.c_str(), dst.c_str(), status_);
EXPECT_TF_OK(status_);
auto result = ReadAll(dst);
EXPECT_TF_OK(status_);
EXPECT_EQ("test_old", result);
}
// These tests below are ported from
// `//tensorflow/core/platform/cloud:gcs_file_system_test`
TEST_F(GCSFilesystemTest, NewRandomAccessFile_NoBlockCache) {
tf_gcs_filesystem::InitTest(filesystem_, false, 0, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(6);
int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
ASSERT_EQ(read, 6) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(6);
int64_t read = tf_random_access_file::Read(file, 0, 6, &result[0], status_);
ASSERT_EQ(read, 6) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "012345") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 6, 6, &result[0], status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "6789") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_ReadAtEOF) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "0123456789", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(10);
int64_t read = tf_random_access_file::Read(file, 0, result.length(),
&result[0], status_);
ASSERT_EQ(read, 10) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "0123456789") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, result.length(), result.length(),
&result[0], status_);
ASSERT_EQ(read, 0) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, NewRandomAccessFile_Buffered_CachedOutOfRange) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
std::string path = GetURIForPath("a_file");
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
ASSERT_TRUE(InsertObject(path, "012345678", &gcs_file->gcs_client, status_));
TF_RandomAccessFile* file = new TF_RandomAccessFile;
tf_gcs_filesystem::NewRandomAccessFile(filesystem_, path.c_str(), file,
status_);
ASSERT_TF_OK(status_);
std::string result;
result.resize(5);
int64_t read = tf_random_access_file::Read(file, 0, result.length(),
&result[0], status_);
ASSERT_EQ(read, 5) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
ASSERT_EQ(result, "01234") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 4, result.length(), &result[0],
status_);
ASSERT_EQ(read, 5) << "Read: " << read << "\n";
ASSERT_TF_OK(status_);
result.resize(read);
ASSERT_EQ(result, "45678") << "Result: " << result << "\n";
read = tf_random_access_file::Read(file, 5, result.length(), &result[0],
status_);
ASSERT_EQ(read, 4) << "Read: " << read << "\n";
ASSERT_EQ(TF_GetCode(status_), TF_OUT_OF_RANGE) << TF_Message(status_);
result.resize(read);
ASSERT_EQ(result, "5678") << "Result: " << result << "\n";
}
TEST_F(GCSFilesystemTest, TestGetStorageClientOptions) {
tf_gcs_filesystem::InitTest(filesystem_, false, 10, 0, 0, 0, 0, status_);
ASSERT_TF_OK(status_) << "Could not initialize filesystem. "
<< TF_Message(status_);
auto gcs_file =
static_cast<tf_gcs_filesystem::GCSFile*>(filesystem_->plugin_filesystem);
// Get the raw connection
auto connection = gcs_file->gcs_client.raw_client();
// Get the options from the connection
auto options = connection->options();
// Access the UserAgentProductsOptions which contain information about
// the UserAgent
auto user_agent_products =
options.get<google::cloud::UserAgentProductsOption>();
bool found_x_goog_api_client = false;
for (const auto& product : user_agent_products) {
if (absl::StrContains(product, "TensorFlow-C-API")) {
found_x_goog_api_client = true;
break;
}
}
EXPECT_TRUE(found_x_goog_api_client);
}
} // namespace
} // namespace tensorflow
GTEST_API_ int main(int argc, char** argv) {
tensorflow::testing::InstallStacktraceHandler();
if (!GetTmpDir()) {
std::cerr << "Could not read GCS_TEST_TMPDIR env";
return -1;
}
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.cc
+42
View File
@@ -0,0 +1,42 @@
/* 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/c/experimental/filesystem/plugins/gcs/gcs_helper.h"
#include <stdio.h>
#include <cstdio>
#include <fstream>
#include <ios>
#include <string>
#include <utility>
TempFile::TempFile(const std::string& temp_file_name, std::ios::openmode mode)
: std::fstream(temp_file_name, mode), name_(temp_file_name) {}
TempFile::TempFile(TempFile&& rhs)
: std::fstream(std::move(rhs)), name_(std::move(rhs.name_)) {}
TempFile::~TempFile() {
std::fstream::close();
std::remove(name_.c_str());
}
const std::string TempFile::getName() const { return name_; }
bool TempFile::truncate() {
std::fstream::close();
std::fstream::open(name_, std::ios::binary | std::ios::out);
return std::fstream::is_open();
}
tensorflow/c/experimental/filesystem/plugins/gcs/gcs_helper.h
+35
View File
@@ -0,0 +1,35 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
#include <fstream>
#include <ios>
#include <string>
class TempFile : public std::fstream {
public:
// We should specify openmode each time we call TempFile.
TempFile(const std::string& temp_file_name, std::ios::openmode mode);
TempFile(TempFile&& rhs);
~TempFile() override;
const std::string getName() const;
bool truncate();
private:
const std::string name_;
};
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_GCS_HELPER_H_
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc
+316
View File
@@ -0,0 +1,316 @@
/* 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/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include <cstdint>
#include <cstring>
#include <iostream>
#include <iterator>
#include <limits>
#include <memory>
#include <sstream>
#include <string>
#include <utility>
#include <vector>
#include "absl/synchronization/mutex.h"
#include "absl/time/time.h"
#include "tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h"
namespace tf_gcs_filesystem {
bool RamFileBlockCache::BlockNotStale(const std::shared_ptr<Block>& block) {
absl::MutexLock l(block->mu);
if (block->state != FetchState::FINISHED) {
return true; // No need to check for staleness.
}
if (max_staleness_ == 0) return true; // Not enforcing staleness.
return timer_seconds_() - block->timestamp <= max_staleness_;
}
std::shared_ptr<RamFileBlockCache::Block> RamFileBlockCache::Lookup(
const Key& key) {
absl::MutexLock lock(mu_);
auto entry = block_map_.find(key);
if (entry != block_map_.end()) {
if (BlockNotStale(entry->second)) {
return entry->second;
} else {
// Remove the stale block and continue.
RemoveFile_Locked(key.first);
}
}
// Insert a new empty block, setting the bookkeeping to sentinel values
// in order to update them as appropriate.
auto new_entry = std::make_shared<Block>();
lru_list_.push_front(key);
lra_list_.push_front(key);
new_entry->lru_iterator = lru_list_.begin();
new_entry->lra_iterator = lra_list_.begin();
new_entry->timestamp = timer_seconds_();
block_map_.emplace(std::make_pair(key, new_entry));
return new_entry;
}
// Remove blocks from the cache until we do not exceed our maximum size.
void RamFileBlockCache::Trim() {
while (!lru_list_.empty() && cache_size_ > max_bytes_) {
RemoveBlock(block_map_.find(lru_list_.back()));
}
}
/// Move the block to the front of the LRU list if it isn't already there.
void RamFileBlockCache::UpdateLRU(const Key& key,
const std::shared_ptr<Block>& block,
TF_Status* status) {
absl::MutexLock lock(mu_);
if (block->timestamp == 0) {
// The block was evicted from another thread. Allow it to remain evicted.
return TF_SetStatus(status, TF_OK, "");
}
if (block->lru_iterator != lru_list_.begin()) {
lru_list_.erase(block->lru_iterator);
lru_list_.push_front(key);
block->lru_iterator = lru_list_.begin();
}
// Check for inconsistent state. If there is a block later in the same file
// in the cache, and our current block is not block size, this likely means
// we have inconsistent state within the cache. Note: it's possible some
// incomplete reads may still go undetected.
if (block->data.size() < block_size_) {
Key fmax = std::make_pair(key.first, std::numeric_limits<size_t>::max());
auto fcmp = block_map_.upper_bound(fmax);
if (fcmp != block_map_.begin() && key < (--fcmp)->first) {
return TF_SetStatus(status, TF_INTERNAL,
"Block cache contents are inconsistent.");
}
}
Trim();
return TF_SetStatus(status, TF_OK, "");
}
void RamFileBlockCache::MaybeFetch(const Key& key,
const std::shared_ptr<Block>& block,
TF_Status* status) {
bool downloaded_block = false;
auto reconcile_state = MakeCleanup([this, &downloaded_block, &key, &block] {
// Perform this action in a cleanup callback to avoid locking mu_ after
// locking block->mu.
if (downloaded_block) {
absl::MutexLock l(mu_);
// Do not update state if the block is already to be evicted.
if (block->timestamp != 0) {
// Use capacity() instead of size() to account for all memory
// used by the cache.
cache_size_ += block->data.capacity();
// Put to beginning of LRA list.
lra_list_.erase(block->lra_iterator);
lra_list_.push_front(key);
block->lra_iterator = lra_list_.begin();
block->timestamp = timer_seconds_();
}
}
});
// Loop until either block content is successfully fetched, or our request
// encounters an error.
absl::MutexLock l(block->mu);
TF_SetStatus(status, TF_OK, "");
while (true) {
switch (block->state) {
case FetchState::ERROR:
// TF_FALLTHROUGH_INTENDED
case FetchState::CREATED:
block->state = FetchState::FETCHING;
block->mu.unlock(); // Release the lock while making the API call.
block->data.clear();
block->data.resize(block_size_, 0);
int64_t bytes_transferred;
bytes_transferred = block_fetcher_(key.first, key.second, block_size_,
block->data.data(), status);
block->mu.lock(); // Reacquire the lock immediately afterwards
if (TF_GetCode(status) == TF_OK) {
block->data.resize(bytes_transferred, 0);
// Shrink the data capacity to the actual size used.
// NOLINTNEXTLINE: shrink_to_fit() may not shrink the capacity.
std::vector<char>(block->data).swap(block->data);
downloaded_block = true;
block->state = FetchState::FINISHED;
} else {
block->state = FetchState::ERROR;
}
block->cond_var.SignalAll();
return;
case FetchState::FETCHING:
block->cond_var.WaitWithTimeout(&block->mu, absl::Minutes(1));
if (block->state == FetchState::FINISHED) {
return TF_SetStatus(status, TF_OK, "");
}
// Re-loop in case of errors.
break;
case FetchState::FINISHED:
return TF_SetStatus(status, TF_OK, "");
}
}
return TF_SetStatus(
status, TF_INTERNAL,
"Control flow should never reach the end of RamFileBlockCache::Fetch.");
}
int64_t RamFileBlockCache::Read(const std::string& filename, size_t offset,
size_t n, char* buffer, TF_Status* status) {
if (n == 0) {
TF_SetStatus(status, TF_OK, "");
return 0;
}
if (!IsCacheEnabled() || (n > max_bytes_)) {
// The cache is effectively disabled, so we pass the read through to the
// fetcher without breaking it up into blocks.
return block_fetcher_(filename, offset, n, buffer, status);
}
// Calculate the block-aligned start and end of the read.
size_t start = block_size_ * (offset / block_size_);
size_t finish = block_size_ * ((offset + n) / block_size_);
if (finish < offset + n) {
finish += block_size_;
}
size_t total_bytes_transferred = 0;
// Now iterate through the blocks, reading them one at a time.
for (size_t pos = start; pos < finish; pos += block_size_) {
Key key = std::make_pair(filename, pos);
// Look up the block, fetching and inserting it if necessary, and update the
// LRU iterator for the key and block.
std::shared_ptr<Block> block = Lookup(key);
if (!block) {
std::cerr << "No block for key " << key.first << "@" << key.second;
abort();
}
MaybeFetch(key, block, status);
if (TF_GetCode(status) != TF_OK) return -1;
UpdateLRU(key, block, status);
if (TF_GetCode(status) != TF_OK) return -1;
// Copy the relevant portion of the block into the result buffer.
const auto& data = block->data;
if (offset >= pos + data.size()) {
// The requested offset is at or beyond the end of the file. This can
// happen if `offset` is not block-aligned, and the read returns the last
// block in the file, which does not extend all the way out to `offset`.
std::stringstream os;
os << "EOF at offset " << offset << " in file " << filename
<< " at position " << pos << " with data size " << data.size();
TF_SetStatus(status, TF_OUT_OF_RANGE, std::move(os).str().c_str());
return total_bytes_transferred;
}
auto begin = data.begin();
if (offset > pos) {
// The block begins before the slice we're reading.
begin += offset - pos;
}
auto end = data.end();
if (pos + data.size() > offset + n) {
// The block extends past the end of the slice we're reading.
end -= (pos + data.size()) - (offset + n);
}
if (begin < end) {
size_t bytes_to_copy = end - begin;
memcpy(&buffer[total_bytes_transferred], &*begin, bytes_to_copy);
total_bytes_transferred += bytes_to_copy;
}
if (data.size() < block_size_) {
// The block was a partial block and thus signals EOF at its upper bound.
break;
}
}
TF_SetStatus(status, TF_OK, "");
return total_bytes_transferred;
}
bool RamFileBlockCache::ValidateAndUpdateFileSignature(
const std::string& filename, int64_t file_signature) {
absl::MutexLock lock(mu_);
auto it = file_signature_map_.find(filename);
if (it != file_signature_map_.end()) {
if (it->second == file_signature) {
return true;
}
// Remove the file from cache if the signatures don't match.
RemoveFile_Locked(filename);
it->second = file_signature;
return false;
}
file_signature_map_[filename] = file_signature;
return true;
}
size_t RamFileBlockCache::CacheSize() const {
absl::MutexLock lock(mu_);
return cache_size_;
}
void RamFileBlockCache::Prune() {
while (!stop_pruning_thread_.WaitForNotificationWithTimeout(
absl::Microseconds(1000000))) {
absl::MutexLock lock(mu_);
uint64_t now = timer_seconds_();
while (!lra_list_.empty()) {
auto it = block_map_.find(lra_list_.back());
if (now - it->second->timestamp <= max_staleness_) {
// The oldest block is not yet expired. Come back later.
break;
}
// We need to make a copy of the filename here, since it could otherwise
// be used within RemoveFile_Locked after `it` is deleted.
RemoveFile_Locked(std::string(it->first.first));
}
}
}
void RamFileBlockCache::Flush() {
absl::MutexLock lock(mu_);
block_map_.clear();
lru_list_.clear();
lra_list_.clear();
cache_size_ = 0;
}
void RamFileBlockCache::RemoveFile(const std::string& filename) {
absl::MutexLock lock(mu_);
RemoveFile_Locked(filename);
}
void RamFileBlockCache::RemoveFile_Locked(const std::string& filename) {
Key begin = std::make_pair(filename, 0);
auto it = block_map_.lower_bound(begin);
while (it != block_map_.end() && it->first.first == filename) {
auto next = std::next(it);
RemoveBlock(it);
it = next;
}
}
void RamFileBlockCache::RemoveBlock(BlockMap::iterator entry) {
// This signals that the block is removed, and should not be inadvertently
// reinserted into the cache in UpdateLRU.
entry->second->timestamp = 0;
lru_list_.erase(entry->second->lru_iterator);
lra_list_.erase(entry->second->lra_iterator);
cache_size_ -= entry->second->data.capacity();
block_map_.erase(entry);
}
} // namespace tf_gcs_filesystem
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h
+270
View File
@@ -0,0 +1,270 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
#include <cstddef>
#include <cstdint>
#include <functional>
#include <iostream>
#include <list>
#include <map>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/base/thread_annotations.h"
#include "absl/log/log.h"
#include "absl/strings/str_format.h"
#include "absl/synchronization/mutex.h"
#include "absl/synchronization/notification.h"
#include "tensorflow/c/env.h"
#include "tensorflow/c/tf_status.h"
namespace tf_gcs_filesystem {
/// \brief An LRU block cache of file contents, keyed by {filename, offset}.
///
/// This class should be shared by read-only random access files on a remote
/// filesystem (e.g. GCS).
class RamFileBlockCache {
public:
/// The callback executed when a block is not found in the cache, and needs to
/// be fetched from the backing filesystem. This callback is provided when the
/// cache is constructed. It returns total bytes read ( -1 in case of errors
/// ). The `status` should be `TF_OK` as long as the read from the remote
/// filesystem succeeded (similar to the semantics of the read(2) system
/// call).
typedef std::function<int64_t(const std::string& filename, size_t offset,
size_t buffer_size, char* buffer,
TF_Status* status)>
BlockFetcher;
RamFileBlockCache(size_t block_size, size_t max_bytes, uint64_t max_staleness,
BlockFetcher block_fetcher,
std::function<uint64_t()> timer_seconds = TF_NowSeconds)
: block_size_(block_size),
max_bytes_(max_bytes),
max_staleness_(max_staleness),
block_fetcher_(block_fetcher),
timer_seconds_(timer_seconds),
pruning_thread_(nullptr,
[](TF_Thread* thread) { TF_JoinThread(thread); }) {
if (max_staleness_ > 0) {
TF_ThreadOptions thread_options;
TF_DefaultThreadOptions(&thread_options);
pruning_thread_.reset(
TF_StartThread(&thread_options, "TF_prune_FBC", PruneThread, this));
}
VLOG(1) << absl::StrFormat("GCS file block cache is %s.\n",
(IsCacheEnabled() ? "enabled" : "disabled"));
}
~RamFileBlockCache() {
if (pruning_thread_) {
stop_pruning_thread_.Notify();
// Destroying pruning_thread_ will block until Prune() receives the above
// notification and returns.
pruning_thread_.reset();
}
}
/// Read `n` bytes from `filename` starting at `offset` into `buffer`. It
/// returns total bytes read ( -1 in case of errors ). This method will set
/// `status` to:
///
/// 1) The error from the remote filesystem, if the read from the remote
/// filesystem failed.
/// 2) `TF_FAILED_PRECONDITION` if the read from the remote filesystem
/// succeeded,
/// but the read returned a partial block, and the LRU cache contained a
/// block at a higher offset (indicating that the partial block should have
/// been a full block).
/// 3) `TF_OUT_OF_RANGE` if the read from the remote filesystem succeeded, but
/// the file contents do not extend past `offset` and thus nothing was
/// placed in `out`.
/// 4) `TF_OK` otherwise (i.e. the read succeeded, and at least one byte was
/// placed
/// in `buffer`).
///
/// Caller is responsible for allocating memory for `buffer`.
/// `buffer` will be left unchanged in case of errors.
int64_t Read(const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status);
// Validate the given file signature with the existing file signature in the
// cache. Returns true if the signature doesn't change or the file doesn't
// exist before. If the signature changes, update the existing signature with
// the new one and remove the file from cache.
bool ValidateAndUpdateFileSignature(const std::string& filename,
int64_t file_signature)
ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all cached blocks for `filename`.
void RemoveFile(const std::string& filename) ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all cached data.
void Flush() ABSL_LOCKS_EXCLUDED(mu_);
/// Accessors for cache parameters.
size_t block_size() const { return block_size_; }
size_t max_bytes() const { return max_bytes_; }
uint64_t max_staleness() const { return max_staleness_; }
/// The current size (in bytes) of the cache.
size_t CacheSize() const ABSL_LOCKS_EXCLUDED(mu_);
// Returns true if the cache is enabled. If false, the BlockFetcher callback
// is always executed during Read.
bool IsCacheEnabled() const { return block_size_ > 0 && max_bytes_ > 0; }
// We can not pass a lambda with capture as a function pointer to
// `TF_StartThread`, so we have to wrap `Prune` inside a static function.
static void PruneThread(void* param) {
auto ram_file_block_cache = static_cast<RamFileBlockCache*>(param);
ram_file_block_cache->Prune();
}
private:
/// The size of the blocks stored in the LRU cache, as well as the size of the
/// reads from the underlying filesystem.
const size_t block_size_;
/// The maximum number of bytes (sum of block sizes) allowed in the LRU cache.
const size_t max_bytes_;
/// The maximum staleness of any block in the LRU cache, in seconds.
const uint64_t max_staleness_;
/// The callback to read a block from the underlying filesystem.
const BlockFetcher block_fetcher_;
/// The callback to read timestamps.
const std::function<uint64_t()> timer_seconds_;
/// \brief The key type for the file block cache.
///
/// The file block cache key is a {filename, offset} pair.
typedef std::pair<std::string, size_t> Key;
/// \brief The state of a block.
///
/// A block begins in the CREATED stage. The first thread will attempt to read
/// the block from the filesystem, transitioning the state of the block to
/// FETCHING. After completing, if the read was successful the state should
/// be FINISHED. Otherwise the state should be ERROR. A subsequent read can
/// re-fetch the block if the state is ERROR.
enum class FetchState {
CREATED,
FETCHING,
FINISHED,
ERROR,
};
/// \brief A block of a file.
///
/// A file block consists of the block data, the block's current position in
/// the LRU cache, the timestamp (seconds since epoch) at which the block
/// was cached, a coordination lock, and state & condition variables.
///
/// Thread safety:
/// The iterator and timestamp fields should only be accessed while holding
/// the block-cache-wide mu_ instance variable. The state variable should only
/// be accessed while holding the Block's mu lock. The data vector should only
/// be accessed after state == FINISHED, and it should never be modified.
///
/// In order to prevent deadlocks, never grab the block-cache-wide mu_ lock
/// AFTER grabbing any block's mu lock. It is safe to grab mu without locking
/// mu_.
struct Block {
/// The block data.
std::vector<char> data;
/// A list iterator pointing to the block's position in the LRU list.
std::list<Key>::iterator lru_iterator;
/// A list iterator pointing to the block's position in the LRA list.
std::list<Key>::iterator lra_iterator;
/// The timestamp (seconds since epoch) at which the block was cached.
uint64_t timestamp;
/// Mutex to guard state variable
absl::Mutex mu;
/// The state of the block.
FetchState state ABSL_GUARDED_BY(mu) = FetchState::CREATED;
/// Wait on cond_var if state is FETCHING.
absl::CondVar cond_var;
};
/// \brief The block map type for the file block cache.
///
/// The block map is an ordered map from Key to Block.
typedef std::map<Key, std::shared_ptr<Block>> BlockMap;
/// Prune the cache by removing files with expired blocks.
void Prune() ABSL_LOCKS_EXCLUDED(mu_);
bool BlockNotStale(const std::shared_ptr<Block>& block)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Look up a Key in the block cache.
std::shared_ptr<Block> Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_);
void MaybeFetch(const Key& key, const std::shared_ptr<Block>& block,
TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
/// Trim the block cache to make room for another entry.
void Trim() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Update the LRU iterator for the block at `key`.
void UpdateLRU(const Key& key, const std::shared_ptr<Block>& block,
TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
/// Remove all blocks of a file, with mu_ already held.
void RemoveFile_Locked(const std::string& filename)
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// Remove the block `entry` from the block map and LRU list, and update the
/// cache size accordingly.
void RemoveBlock(BlockMap::iterator entry) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
/// The cache pruning thread that removes files with expired blocks.
std::unique_ptr<TF_Thread, std::function<void(TF_Thread*)>> pruning_thread_;
/// Notification for stopping the cache pruning thread.
absl::Notification stop_pruning_thread_;
/// Guards access to the block map, LRU list, and cached byte count.
mutable absl::Mutex mu_;
/// The block map (map from Key to Block).
BlockMap block_map_ ABSL_GUARDED_BY(mu_);
/// The LRU list of block keys. The front of the list identifies the most
/// recently accessed block.
std::list<Key> lru_list_ ABSL_GUARDED_BY(mu_);
/// The LRA (least recently added) list of block keys. The front of the list
/// identifies the most recently added block.
///
/// Note: blocks are added to lra_list_ only after they have successfully been
/// fetched from the underlying block store.
std::list<Key> lra_list_ ABSL_GUARDED_BY(mu_);
/// The combined number of bytes in all of the cached blocks.
size_t cache_size_ ABSL_GUARDED_BY(mu_) = 0;
// A filename->file_signature map.
std::map<std::string, int64_t> file_signature_map_ ABSL_GUARDED_BY(mu_);
};
} // namespace tf_gcs_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_GCS_RAM_FILE_BLOCK_CACHE_H_
tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc
+614
View File
@@ -0,0 +1,614 @@
/* 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/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h"
#include <algorithm>
#include <cstdint>
#include <cstring>
#include <list>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/ascii.h"
#include "absl/synchronization/blocking_counter.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/cloud/now_seconds_env.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace {
absl::Status ReadCache(tf_gcs_filesystem::RamFileBlockCache* cache,
const std::string& filename, size_t offset, size_t n,
std::vector<char>* out) {
out->clear();
out->resize(n, 0);
TF_Status status;
auto bytes_transferred =
cache->Read(filename, offset, n, out->data(), &status);
if (bytes_transferred >= 0) {
EXPECT_LE(bytes_transferred, n);
out->resize(bytes_transferred, n);
}
return status.status;
}
TEST(RamFileBlockCacheTest, IsCacheEnabled) {
auto fetcher = [](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
// Do nothing.
TF_SetStatus(status, TF_OK, "");
return 0;
};
tf_gcs_filesystem::RamFileBlockCache cache1(0, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache2(16, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache3(0, 32, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache4(16, 32, 0, fetcher);
EXPECT_FALSE(cache1.IsCacheEnabled());
EXPECT_FALSE(cache2.IsCacheEnabled());
EXPECT_FALSE(cache3.IsCacheEnabled());
EXPECT_TRUE(cache4.IsCacheEnabled());
}
TEST(RamFileBlockCacheTest, ValidateAndUpdateFileSignature) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::string filename = "file";
tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
std::vector<char> out;
// First read.
EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 1);
// Second read. Hit cache.
EXPECT_TRUE(cache.ValidateAndUpdateFileSignature(filename, 123));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 1);
// Third read. File signatures are different.
EXPECT_FALSE(cache.ValidateAndUpdateFileSignature(filename, 321));
TF_EXPECT_OK(ReadCache(&cache, filename, 0, 16, &out));
EXPECT_EQ(calls, 2);
}
TEST(RamFileBlockCacheTest, PassThrough) {
const std::string want_filename = "foo/bar";
const size_t want_offset = 42;
const size_t want_n = 1024;
int calls = 0;
auto fetcher = [&calls, want_filename, want_offset, want_n](
const std::string& got_filename, size_t got_offset,
size_t got_n, char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(got_filename, want_filename);
EXPECT_EQ(got_offset, want_offset);
EXPECT_EQ(got_n, want_n);
calls++;
memset(buffer, 'x', got_n);
TF_SetStatus(status, TF_OK, "");
return got_n;
};
// If block_size, max_bytes, or both are zero, or want_n is larger than
// max_bytes the cache is a pass-through.
tf_gcs_filesystem::RamFileBlockCache cache1(1, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache2(0, 1, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache3(0, 0, 0, fetcher);
tf_gcs_filesystem::RamFileBlockCache cache4(1000, 1000, 0, fetcher);
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache1, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 1);
TF_EXPECT_OK(ReadCache(&cache2, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 2);
TF_EXPECT_OK(ReadCache(&cache3, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache4, want_filename, want_offset, want_n, &out));
EXPECT_EQ(calls, 4);
}
TEST(RamFileBlockCacheTest, BlockAlignment) {
// Initialize a 256-byte buffer. This is the file underlying the reads we'll
// do in this test.
const size_t size = 256;
std::vector<char> buf;
buf.reserve(size);
for (int i = 0; i < size; i++) {
buf.push_back(i);
}
// The fetcher just fetches slices of the buffer.
auto fetcher = [&buf](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
int64_t bytes_transferred;
if (offset < buf.size()) {
size_t bytes_to_copy = std::min<size_t>(buf.size() - offset, n);
memcpy(buffer, buf.data() + offset, bytes_to_copy);
bytes_transferred = bytes_to_copy;
} else {
bytes_transferred = 0;
}
TF_SetStatus(status, TF_OK, "");
return bytes_transferred;
};
for (size_t block_size = 2; block_size <= 4; block_size++) {
// Make a cache of N-byte block size (1 block) and verify that reads of
// varying offsets and lengths return correct data.
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
for (size_t offset = 0; offset < 10; offset++) {
for (size_t n = block_size - 2; n <= block_size + 2; n++) {
std::vector<char> got;
TF_EXPECT_OK(ReadCache(&cache, "", offset, n, &got));
// Verify the size of the read.
if (offset + n <= size) {
// Expect a full read.
EXPECT_EQ(got.size(), n) << "block size = " << block_size
<< ", offset = " << offset << ", n = " << n;
} else {
// Expect a partial read.
EXPECT_EQ(got.size(), size - offset)
<< "block size = " << block_size << ", offset = " << offset
<< ", n = " << n;
}
// Verify the contents of the read.
std::vector<char>::const_iterator begin = buf.begin() + offset;
std::vector<char>::const_iterator end =
offset + n > buf.size() ? buf.end() : begin + n;
std::vector<char> want(begin, end);
EXPECT_EQ(got, want) << "block size = " << block_size
<< ", offset = " << offset << ", n = " << n;
}
}
}
}
TEST(RamFileBlockCacheTest, CacheHits) {
const size_t block_size = 16;
std::set<size_t> calls;
auto fetcher = [&calls, block_size](const std::string& filename,
size_t offset, size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
EXPECT_EQ(calls.find(offset), calls.end()) << "at offset " << offset;
calls.insert(offset);
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const uint32_t block_count = 256;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, block_count * block_size, 0, fetcher);
std::vector<char> out;
out.resize(block_count, 0);
// The cache has space for `block_count` blocks. The loop with i = 0 should
// fill the cache, and the loop with i = 1 should be all cache hits. The
// fetcher checks that it is called once and only once for each offset (to
// fetch the corresponding block).
for (int i = 0; i < 2; i++) {
for (int j = 0; j < block_count; j++) {
TF_EXPECT_OK(ReadCache(&cache, "", block_size * j, block_size, &out));
}
}
}
TEST(RamFileBlockCacheTest, OutOfRange) {
// Tests reads of a 24-byte file with block size 16.
const size_t block_size = 16;
const size_t file_size = 24;
bool first_block = false;
bool second_block = false;
auto fetcher = [block_size, file_size, &first_block, &second_block](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
size_t bytes_to_copy = 0;
if (offset == 0) {
// The first block (16 bytes) of the file.
memset(buffer, 'x', n);
bytes_to_copy = n;
first_block = true;
} else if (offset == block_size) {
// The second block (8 bytes) of the file.
bytes_to_copy = file_size - block_size;
memset(buffer, 'x', bytes_to_copy);
second_block = true;
}
TF_SetStatus(status, TF_OK, "");
return bytes_to_copy;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
std::vector<char> out;
// Reading the first 16 bytes should be fine.
TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size, &out));
EXPECT_TRUE(first_block);
EXPECT_EQ(out.size(), block_size);
// Reading at offset file_size + 4 will read the second block (since the read
// at file_size + 4 = 28 will be aligned to an offset of 16) but will return
// OutOfRange because the offset is past the end of the 24-byte file.
absl::Status status = ReadCache(&cache, "", file_size + 4, 4, &out);
EXPECT_EQ(status.code(), error::OUT_OF_RANGE);
EXPECT_TRUE(second_block);
// Reading the second full block will return 8 bytes, from a cache hit.
second_block = false;
TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
EXPECT_FALSE(second_block);
EXPECT_EQ(out.size(), file_size - block_size);
}
TEST(RamFileBlockCacheTest, Inconsistent) {
// Tests the detection of interrupted reads leading to partially filled blocks
// where we expected complete blocks.
const size_t block_size = 16;
// This fetcher returns OK but only fills in one byte for any offset.
auto fetcher = [block_size](const std::string& filename, size_t offset,
size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset % block_size, 0);
EXPECT_GE(n, 1);
memset(buffer, 'x', 1);
TF_SetStatus(status, TF_OK, "");
return 1;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, 2 * block_size, 0,
fetcher);
std::vector<char> out;
// Read the second block; this should yield an OK status and a single byte.
TF_EXPECT_OK(ReadCache(&cache, "", block_size, block_size, &out));
EXPECT_EQ(out.size(), 1);
// Now read the first block; this should yield an INTERNAL error because we
// had already cached a partial block at a later position.
absl::Status status = ReadCache(&cache, "", 0, block_size, &out);
EXPECT_EQ(status.code(), error::INTERNAL);
}
TEST(RamFileBlockCacheTest, LRU) {
const size_t block_size = 16;
std::list<size_t> calls;
auto fetcher = [&calls, block_size](const std::string& filename,
size_t offset, size_t n, char* buffer,
TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_FALSE(calls.empty()) << "at offset = " << offset;
if (!calls.empty()) {
EXPECT_EQ(offset, calls.front());
calls.pop_front();
}
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const uint32_t block_count = 2;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, block_count * block_size, 0, fetcher);
std::vector<char> out;
// Read blocks from the cache, and verify the LRU behavior based on the
// fetcher calls that the cache makes.
calls.push_back(0);
// Cache miss - drains an element from `calls`.
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
// Cache hit - does not drain an element from `calls`.
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
calls.push_back(block_size);
// Cache miss followed by cache hit.
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
calls.push_back(2 * block_size);
// Cache miss followed by cache hit. Causes eviction of LRU element.
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
// LRU element was at offset 0. Cache miss.
calls.push_back(0);
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
// Element at 2 * block_size is still in cache, and this read should update
// its position in the LRU list so it doesn't get evicted by the next read.
TF_EXPECT_OK(ReadCache(&cache, "", 2 * block_size, 1, &out));
// Element at block_size was evicted. Reading this element will also cause
// the LRU element (at 0) to be evicted.
calls.push_back(block_size);
TF_EXPECT_OK(ReadCache(&cache, "", block_size, 1, &out));
// Element at 0 was evicted again.
calls.push_back(0);
TF_EXPECT_OK(ReadCache(&cache, "", 0, 1, &out));
}
TEST(RamFileBlockCacheTest, MaxStaleness) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::vector<char> out;
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
// Create a cache with max staleness of 2 seconds, and verify that it works as
// expected.
tf_gcs_filesystem::RamFileBlockCache cache1(
8, 16, 2 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Execute the first read to load the block.
TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
// Now advance the clock one second at a time and redo the read. The call
// count should advance every 3 seconds (i.e. every time the staleness is
// greater than 2).
for (int i = 1; i <= 10; i++) {
env->SetNowSeconds(i + 1);
TF_EXPECT_OK(ReadCache(&cache1, "", 0, 1, &out));
EXPECT_EQ(calls, 1 + i / 3);
}
// Now create a cache with max staleness of 0, and verify that it also works
// as expected.
calls = 0;
env->SetNowSeconds(0);
tf_gcs_filesystem::RamFileBlockCache cache2(
8, 16, 0 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Execute the first read to load the block.
TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
// Advance the clock by a huge amount and verify that the cached block is
// used to satisfy the read.
env->SetNowSeconds(365 * 24 * 60 * 60); // ~1 year, just for fun.
TF_EXPECT_OK(ReadCache(&cache2, "", 0, 1, &out));
EXPECT_EQ(calls, 1);
}
TEST(RamFileBlockCacheTest, RemoveFile) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
char c = (filename == "a") ? 'a' : (filename == "b") ? 'b' : 'x';
if (offset > 0) {
// The first block is lower case and all subsequent blocks are upper case.
c = absl::ascii_toupper(c);
}
memset(buffer, c, n);
TF_SetStatus(status, TF_OK, "");
return n;
};
// This cache has space for 4 blocks; we'll read from two files.
const size_t n = 3;
tf_gcs_filesystem::RamFileBlockCache cache(8, 32, 0, fetcher);
std::vector<char> out;
std::vector<char> a(n, 'a');
std::vector<char> b(n, 'b');
std::vector<char> A(n, 'A');
std::vector<char> B(n, 'B');
// Fill the cache.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
EXPECT_EQ(calls, 1);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
EXPECT_EQ(calls, 2);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// All four blocks should be in the cache now.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// Remove the blocks from "a".
cache.RemoveFile("a");
// Both blocks from "b" should still be there.
TF_EXPECT_OK(ReadCache(&cache, "b", 0, n, &out));
EXPECT_EQ(out, b);
TF_EXPECT_OK(ReadCache(&cache, "b", 8, n, &out));
EXPECT_EQ(out, B);
EXPECT_EQ(calls, 4);
// The blocks from "a" should not be there.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, n, &out));
EXPECT_EQ(out, a);
EXPECT_EQ(calls, 5);
TF_EXPECT_OK(ReadCache(&cache, "a", 8, n, &out));
EXPECT_EQ(out, A);
EXPECT_EQ(calls, 6);
}
TEST(RamFileBlockCacheTest, Prune) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
std::vector<char> out;
// Our fake environment is initialized with the current timestamp.
std::unique_ptr<NowSecondsEnv> env(new NowSecondsEnv);
uint64_t now = Env::Default()->NowSeconds();
env->SetNowSeconds(now);
tf_gcs_filesystem::RamFileBlockCache cache(
8, 32, 1 /* max staleness */, fetcher,
[&env]() { return env->NowSeconds(); });
// Read three blocks into the cache, and advance the timestamp by one second
// with each read. Start with a block of "a" at the current timestamp `now`.
TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
// Now load a block of a different file "b" at timestamp `now` + 1
env->SetNowSeconds(now + 1);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
// Now load a different block of file "a" at timestamp `now` + 1. When the
// first block of "a" expires, this block should also be removed because it
// also belongs to file "a".
TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
// Ensure that all blocks are in the cache (i.e. reads are cache hits).
EXPECT_EQ(cache.CacheSize(), 24);
EXPECT_EQ(calls, 3);
TF_EXPECT_OK(ReadCache(&cache, "a", 0, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
TF_EXPECT_OK(ReadCache(&cache, "a", 8, 1, &out));
EXPECT_EQ(calls, 3);
// Advance the fake timestamp so that "a" becomes stale via its first block.
env->SetNowSeconds(now + 2);
// The pruning thread periodically compares env->NowSeconds() with the oldest
// block's timestamp to see if it should evict any files. At the current fake
// timestamp of `now` + 2, file "a" is stale because its first block is stale,
// but file "b" is not stale yet. Thus, once the pruning thread wakes up (in
// one second of wall time), it should remove "a" and leave "b" alone.
uint64_t start = Env::Default()->NowSeconds();
do {
Env::Default()->SleepForMicroseconds(100000);
} while (cache.CacheSize() == 24 && Env::Default()->NowSeconds() - start < 3);
// There should be one block left in the cache, and it should be the first
// block of "b".
EXPECT_EQ(cache.CacheSize(), 8);
TF_EXPECT_OK(ReadCache(&cache, "b", 0, 1, &out));
EXPECT_EQ(calls, 3);
// Advance the fake time to `now` + 3, at which point "b" becomes stale.
env->SetNowSeconds(now + 3);
// Wait for the pruner to remove "b".
start = Env::Default()->NowSeconds();
do {
Env::Default()->SleepForMicroseconds(100000);
} while (cache.CacheSize() == 8 && Env::Default()->NowSeconds() - start < 3);
// The cache should now be empty.
EXPECT_EQ(cache.CacheSize(), 0);
}
TEST(RamFileBlockCacheTest, ParallelReads) {
// This fetcher won't respond until either `callers` threads are calling it
// concurrently (at which point it will respond with success to all callers),
// or 10 seconds have elapsed (at which point it will respond with an error).
const int callers = 4;
absl::BlockingCounter counter(callers);
absl::Notification notification;
auto fetcher = [&counter, &notification](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
if (counter.DecrementCount()) {
notification.Notify();
// This call to `Wait()` is not expected to block. Calling `Wait()` here
// allows us to satisfy `BlockingCounter`'s requirement: "When `Wait()`
// returns, it is legal to destroy the `BlockingCounter`.".
counter.Wait();
}
if (!notification.WaitForNotificationWithTimeout(absl::Seconds(10))) {
// This avoids having the test time out, which is harder to debug.
TF_SetStatus(status, TF_FAILED_PRECONDITION,
"desired concurrency not reached");
return -1;
}
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
const int block_size = 8;
tf_gcs_filesystem::RamFileBlockCache cache(
block_size, 2 * callers * block_size, 0, fetcher);
std::vector<std::unique_ptr<Thread>> threads;
threads.reserve(callers);
for (int i = 0; i < callers; i++) {
threads.emplace_back(
Env::Default()->StartThread({}, "caller", [block_size, &cache, i]() {
std::vector<char> out;
TF_EXPECT_OK(
ReadCache(&cache, "a", i * block_size, block_size, &out));
std::vector<char> x(block_size, 'x');
EXPECT_EQ(out, x);
}));
}
// The `threads` destructor blocks until the threads can be joined, once their
// respective reads finish (which happens once they are all concurrently being
// executed, or 10 seconds have passed).
}
TEST(RamFileBlockCacheTest, CoalesceConcurrentReads) {
// Concurrent reads to the same file blocks should be de-duplicated.
const size_t block_size = 16;
int num_requests = 0;
absl::Notification notification;
auto fetcher = [&num_requests, &notification, block_size](
const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
EXPECT_EQ(n, block_size);
EXPECT_EQ(offset, 0);
num_requests++;
memset(buffer, 'x', n);
notification.Notify();
// Wait for other thread to issue read.
Env::Default()->SleepForMicroseconds(100000); // 0.1 secs
TF_SetStatus(status, TF_OK, "");
return n;
};
tf_gcs_filesystem::RamFileBlockCache cache(block_size, block_size, 0,
fetcher);
// Fork off thread for parallel read.
std::unique_ptr<Thread> concurrent(
Env::Default()->StartThread({}, "concurrent", [block_size, &cache] {
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", 0, block_size / 2, &out));
EXPECT_EQ(out.size(), block_size / 2);
}));
notification.WaitForNotification();
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", block_size / 2, block_size / 2, &out));
EXPECT_EQ(out.size(), block_size / 2);
EXPECT_EQ(1, num_requests);
}
TEST(RamFileBlockCacheTest, Flush) {
int calls = 0;
auto fetcher = [&calls](const std::string& filename, size_t offset, size_t n,
char* buffer, TF_Status* status) -> int64_t {
calls++;
memset(buffer, 'x', n);
TF_SetStatus(status, TF_OK, "");
return n;
};
tf_gcs_filesystem::RamFileBlockCache cache(16, 32, 0, fetcher);
std::vector<char> out;
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
EXPECT_EQ(calls, 1);
cache.Flush();
TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
EXPECT_EQ(calls, 2);
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/posix/BUILD
+67
View File
@@ -0,0 +1,67 @@
# Experimental posix filesystem plugin.
load("//tensorflow:tensorflow.bzl", "tf_cc_shared_object")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
# Filesystem implementation for POSIX environments: Linux, MacOS, Android, etc.
tf_cc_shared_object(
name = "libposix_filesystem.so",
framework_so = [],
linkstatic = False,
visibility = ["//visibility:public"],
deps = [":posix_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "posix_filesystem_impl",
srcs = ["posix_filesystem.cc"],
hdrs = ["posix_filesystem.h"],
deps = [
":posix_filesystem_helper",
"//tensorflow/c:tf_file_statistics",
"//tensorflow/c:tf_status",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
],
)
# Since building pip package and API tests require a filesystem, we provide a
# static registration target that they should link against.
cc_library(
name = "posix_filesystem_static",
srcs = ["posix_filesystem_static.cc"],
visibility = ["//visibility:public"],
deps = [
":posix_filesystem_impl",
"//tensorflow/c/experimental/filesystem:filesystem_interface",
"//tensorflow/c/experimental/filesystem:modular_filesystem",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/log",
],
alwayslink = 1,
)
# Library implementing helper functionality, so that the above only contains
# the API implementation for modular filesystems.
cc_library(
name = "posix_filesystem_helper",
srcs = ["posix_filesystem_helper.cc"],
hdrs = ["posix_filesystem_helper.h"],
deps = [":copy_file"],
)
# On Linux, we can copy files faster using `sendfile`. But not elsewhere.
# Hence, this private library to select which implementation to use.
cc_library(
name = "copy_file",
srcs = select({
"//tensorflow:linux_x86_64": ["copy_file_linux.cc"],
"//conditions:default": ["copy_file_portable.cc"],
}),
hdrs = ["copy_file.h"],
)
tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h
+32
View File
@@ -0,0 +1,32 @@
/* Copyright 2019 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_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
#include <sys/stat.h>
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses `sendfile` if available (i.e., linux 2.6.33 or later) or an
// intermediate buffer if not.
//
// Returns number of bytes transferred or -1 on failure.
int CopyFileContents(int dst_fd, int src_fd, off_t size);
} // namespace tf_posix_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_COPY_FILE_H_
tensorflow/c/experimental/filesystem/plugins/posix/copy_file_linux.cc
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2019 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 <limits.h>
#include <stdint.h>
#include <sys/sendfile.h>
#include <cstddef>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses `sendfile` specific to linux after 2.6.33.
int CopyFileContents(int dst_fd, int src_fd, off_t size) {
off_t offset = 0;
int bytes_transferred = 0;
int rc = 1;
// When `sendfile` returns 0 we stop copying and let callers handle this.
while (offset < size && rc > 0) {
// Use uint64 for safe compare SSIZE_MAX
uint64_t chunk = size - offset;
if (chunk > SSIZE_MAX) chunk = SSIZE_MAX;
rc = sendfile(dst_fd, src_fd, &offset, static_cast<size_t>(chunk));
if (rc < 0) return -1;
bytes_transferred += rc;
}
return bytes_transferred;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/copy_file_portable.cc
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2019 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 <stdint.h>
#include <unistd.h>
#include <memory>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
// Transfers up to `size` bytes from `dst_fd` to `src_fd`.
//
// This method uses a temporary buffer to hold contents.
int CopyFileContents(int dst_fd, int src_fd, off_t size) {
// Use a copy buffer of 128KB but don't store it on the stack.
constexpr static size_t kPosixCopyFileBufferSize = 128 * 1024;
std::unique_ptr<char[]> buffer(new char[kPosixCopyFileBufferSize]);
off_t offset = 0;
int bytes_transferred = 0;
int rc = 1;
// When `sendfile` returns 0 we stop copying and let callers handle this.
while (offset < size && rc > 0) {
size_t chunk = size - offset;
if (chunk > kPosixCopyFileBufferSize) chunk = kPosixCopyFileBufferSize;
rc = read(src_fd, buffer.get(), chunk);
if (rc < 0) return -1;
int total_write = 0;
int total_read = rc;
while (total_write < total_read && rc > 0) {
rc = write(dst_fd, buffer.get() + total_write, total_read - total_write);
if (rc < 0) return -1;
total_write += rc;
bytes_transferred += rc;
offset += rc;
}
}
return bytes_transferred;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.cc
+461
View File
@@ -0,0 +1,461 @@
/* Copyright 2019 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/c/experimental/filesystem/plugins/posix/posix_filesystem.h"
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h"
#include "tensorflow/c/tf_file_statistics.h"
#include "tensorflow/c/tf_status.h"
// Implementation of a filesystem for POSIX environments.
// This filesystem will support `file://` and empty (local) URI schemes.
static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
static void plugin_memory_free(void* ptr) { free(ptr); }
// SECTION 1. Implementation for `TF_RandomAccessFile`
// ----------------------------------------------------------------------------
namespace tf_random_access_file {
typedef struct PosixFile {
const char* filename;
int fd;
} PosixFile;
static void Cleanup(TF_RandomAccessFile* file) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
close(posix_file->fd);
// This would be safe to free using `free` directly as it is only opaque.
// However, it is better to be consistent everywhere.
plugin_memory_free(const_cast<char*>(posix_file->filename));
delete posix_file;
}
static int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
char* buffer, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
char* dst = buffer;
int64_t read = 0;
while (n > 0) {
// Some platforms, notably macs, throw `EINVAL` if `pread` is asked to read
// more than fits in a 32-bit integer.
size_t requested_read_length;
if (n > INT32_MAX)
requested_read_length = INT32_MAX;
else
requested_read_length = n;
// `pread` returns a `ssize_t` on POSIX, but due to interface being
// cross-platform, return type of `Read` is `int64_t`.
int64_t r = int64_t{pread(posix_file->fd, dst, requested_read_length,
static_cast<off_t>(offset))};
if (r > 0) {
dst += r;
offset += static_cast<uint64_t>(r);
n -= r; // safe as 0 < r <= n so n will never underflow
read += r;
} else if (r == 0) {
TF_SetStatus(status, TF_OUT_OF_RANGE, "Read fewer bytes than requested");
break;
} else if (errno == EINTR || errno == EAGAIN) {
// Retry
} else {
TF_SetStatusFromIOError(status, errno, posix_file->filename);
break;
}
}
return read;
}
} // namespace tf_random_access_file
// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {
typedef struct PosixFile {
const char* filename;
FILE* handle;
} PosixFile;
static void Cleanup(TF_WritableFile* file) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
plugin_memory_free(const_cast<char*>(posix_file->filename));
delete posix_file;
}
static void Append(const TF_WritableFile* file, const char* buffer, size_t n,
TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
size_t r = fwrite(buffer, 1, n, posix_file->handle);
if (r != n)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
}
static int64_t Tell(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
// POSIX's `ftell` returns `long`, do a manual cast.
int64_t position = int64_t{ftell(posix_file->handle)};
if (position < 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
return position;
}
static void Flush(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
TF_SetStatus(status, TF_OK, "");
if (fflush(posix_file->handle) != 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
}
static void Sync(const TF_WritableFile* file, TF_Status* status) {
// For historical reasons, this does the same as `Flush` at the moment.
// TODO(b/144055243): This should use `fsync`/`sync`.
Flush(file, status);
}
static void Close(const TF_WritableFile* file, TF_Status* status) {
auto posix_file = static_cast<PosixFile*>(file->plugin_file);
if (fclose(posix_file->handle) != 0)
TF_SetStatusFromIOError(status, errno, posix_file->filename);
else
TF_SetStatus(status, TF_OK, "");
}
} // namespace tf_writable_file
// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {
typedef struct PosixMemoryRegion {
const void* const address;
const uint64_t length;
} PosixMemoryRegion;
static void Cleanup(TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
munmap(const_cast<void*>(r->address), r->length);
delete r;
}
static const void* Data(const TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
return r->address;
}
static uint64_t Length(const TF_ReadOnlyMemoryRegion* region) {
auto r = static_cast<PosixMemoryRegion*>(region->plugin_memory_region);
return r->length;
}
} // namespace tf_read_only_memory_region
// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_posix_filesystem {
static void Init(TF_Filesystem* filesystem, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
}
static void Cleanup(TF_Filesystem* filesystem) {}
static void NewRandomAccessFile(const TF_Filesystem* filesystem,
const char* path, TF_RandomAccessFile* file,
TF_Status* status) {
int fd = open(path, O_RDONLY);
if (fd < 0) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
struct stat st;
fstat(fd, &st);
if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "path is a directory");
close(fd);
return;
}
file->plugin_file = new tf_random_access_file::PosixFile({strdup(path), fd});
TF_SetStatus(status, TF_OK, "");
}
static void NewWritableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status) {
FILE* f = fopen(path, "w");
if (f == nullptr) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
file->plugin_file = new tf_writable_file::PosixFile({strdup(path), f});
TF_SetStatus(status, TF_OK, "");
}
static void NewAppendableFile(const TF_Filesystem* filesystem, const char* path,
TF_WritableFile* file, TF_Status* status) {
FILE* f = fopen(path, "a");
if (f == nullptr) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
file->plugin_file = new tf_writable_file::PosixFile({strdup(path), f});
TF_SetStatus(status, TF_OK, "");
}
static void NewReadOnlyMemoryRegionFromFile(const TF_Filesystem* filesystem,
const char* path,
TF_ReadOnlyMemoryRegion* region,
TF_Status* status) {
int fd = open(path, O_RDONLY);
if (fd < 0) {
TF_SetStatusFromIOError(status, errno, path);
return;
}
struct stat st;
fstat(fd, &st);
if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "path is a directory");
} else {
const void* address =
mmap(nullptr, st.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
if (address == MAP_FAILED) {
TF_SetStatusFromIOError(status, errno, path);
} else {
region->plugin_memory_region =
new tf_read_only_memory_region::PosixMemoryRegion{
address, static_cast<uint64_t>(st.st_size)};
TF_SetStatus(status, TF_OK, "");
}
}
close(fd);
}
static void CreateDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (strlen(path) == 0)
TF_SetStatus(status, TF_ALREADY_EXISTS, "already exists");
else if (mkdir(path, /*mode=*/0755) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void DeleteFile(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (unlink(path) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void DeleteDir(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (rmdir(path) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void RenameFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status) {
// If target is a directory return TF_FAILED_PRECONDITION.
// Target might be missing, so don't error in that case.
struct stat st;
if (stat(dst, &st) != 0) {
if (errno != ENOENT) {
TF_SetStatusFromIOError(status, errno, dst);
return;
}
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "target path is a directory");
return;
}
// We cannot rename directories yet, so prevent this.
if (stat(src, &st) != 0) {
TF_SetStatusFromIOError(status, errno, src);
return;
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "source path is a directory");
return;
}
// Do the actual rename. Here both arguments are filenames.
if (rename(src, dst) != 0)
TF_SetStatusFromIOError(status, errno, dst);
else
TF_SetStatus(status, TF_OK, "");
}
static void CopyFile(const TF_Filesystem* filesystem, const char* src,
const char* dst, TF_Status* status) {
// If target is a directory return TF_FAILED_PRECONDITION.
// Target might be missing, so don't error in that case.
struct stat st;
if (stat(dst, &st) != 0) {
if (errno != ENOENT) {
TF_SetStatusFromIOError(status, errno, dst);
return;
}
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "target path is a directory");
return;
}
// We cannot copy directories yet, so prevent this.
if (stat(src, &st) != 0) {
TF_SetStatusFromIOError(status, errno, src);
return;
} else if (S_ISDIR(st.st_mode)) {
TF_SetStatus(status, TF_FAILED_PRECONDITION, "source path is a directory");
return;
}
// Both `src` and `dst` point to files here. Delegate to helper.
if (TransferFileContents(src, dst, st.st_mode, st.st_size) < 0)
TF_SetStatusFromIOError(status, errno, dst);
else
TF_SetStatus(status, TF_OK, "");
}
static void PathExists(const TF_Filesystem* filesystem, const char* path,
TF_Status* status) {
if (access(path, F_OK) != 0)
TF_SetStatusFromIOError(status, errno, path);
else
TF_SetStatus(status, TF_OK, "");
}
static void Stat(const TF_Filesystem* filesystem, const char* path,
TF_FileStatistics* stats, TF_Status* status) {
struct stat sbuf;
if (stat(path, &sbuf) != 0) {
TF_SetStatusFromIOError(status, errno, path);
} else {
stats->length = sbuf.st_size;
stats->mtime_nsec = sbuf.st_mtime * (1000 * 1000 * 1000);
stats->is_directory = S_ISDIR(sbuf.st_mode);
TF_SetStatus(status, TF_OK, "");
}
}
static int GetChildren(const TF_Filesystem* filesystem, const char* path,
char*** entries, TF_Status* status) {
struct dirent** dir_entries = nullptr;
/* we don't promise entries would be sorted */
int num_entries =
scandir(path, &dir_entries, RemoveSpecialDirectoryEntries, nullptr);
if (num_entries < 0) {
TF_SetStatusFromIOError(status, errno, path);
} else {
*entries = static_cast<char**>(
plugin_memory_allocate(num_entries * sizeof((*entries)[0])));
for (int i = 0; i < num_entries; i++) {
(*entries)[i] = strdup(dir_entries[i]->d_name);
plugin_memory_free(dir_entries[i]);
}
plugin_memory_free(dir_entries);
}
return num_entries;
}
} // namespace tf_posix_filesystem
static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
const char* uri) {
TF_SetFilesystemVersionMetadata(ops);
ops->scheme = strdup(uri);
ops->random_access_file_ops = static_cast<TF_RandomAccessFileOps*>(
plugin_memory_allocate(TF_RANDOM_ACCESS_FILE_OPS_SIZE));
ops->random_access_file_ops->cleanup = tf_random_access_file::Cleanup;
ops->random_access_file_ops->read = tf_random_access_file::Read;
ops->writable_file_ops = static_cast<TF_WritableFileOps*>(
plugin_memory_allocate(TF_WRITABLE_FILE_OPS_SIZE));
ops->writable_file_ops->cleanup = tf_writable_file::Cleanup;
ops->writable_file_ops->append = tf_writable_file::Append;
ops->writable_file_ops->tell = tf_writable_file::Tell;
ops->writable_file_ops->flush = tf_writable_file::Flush;
ops->writable_file_ops->sync = tf_writable_file::Sync;
ops->writable_file_ops->close = tf_writable_file::Close;
ops->read_only_memory_region_ops = static_cast<TF_ReadOnlyMemoryRegionOps*>(
plugin_memory_allocate(TF_READ_ONLY_MEMORY_REGION_OPS_SIZE));
ops->read_only_memory_region_ops->cleanup =
tf_read_only_memory_region::Cleanup;
ops->read_only_memory_region_ops->data = tf_read_only_memory_region::Data;
ops->read_only_memory_region_ops->length = tf_read_only_memory_region::Length;
ops->filesystem_ops = static_cast<TF_FilesystemOps*>(
plugin_memory_allocate(TF_FILESYSTEM_OPS_SIZE));
ops->filesystem_ops->init = tf_posix_filesystem::Init;
ops->filesystem_ops->cleanup = tf_posix_filesystem::Cleanup;
ops->filesystem_ops->new_random_access_file =
tf_posix_filesystem::NewRandomAccessFile;
ops->filesystem_ops->new_writable_file = tf_posix_filesystem::NewWritableFile;
ops->filesystem_ops->new_appendable_file =
tf_posix_filesystem::NewAppendableFile;
ops->filesystem_ops->new_read_only_memory_region_from_file =
tf_posix_filesystem::NewReadOnlyMemoryRegionFromFile;
ops->filesystem_ops->create_dir = tf_posix_filesystem::CreateDir;
ops->filesystem_ops->delete_file = tf_posix_filesystem::DeleteFile;
ops->filesystem_ops->delete_dir = tf_posix_filesystem::DeleteDir;
ops->filesystem_ops->rename_file = tf_posix_filesystem::RenameFile;
ops->filesystem_ops->copy_file = tf_posix_filesystem::CopyFile;
ops->filesystem_ops->path_exists = tf_posix_filesystem::PathExists;
ops->filesystem_ops->stat = tf_posix_filesystem::Stat;
ops->filesystem_ops->get_children = tf_posix_filesystem::GetChildren;
}
void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
info->plugin_memory_allocate = plugin_memory_allocate;
info->plugin_memory_free = plugin_memory_free;
info->num_schemes = 2;
info->ops = static_cast<TF_FilesystemPluginOps*>(
plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
ProvideFilesystemSupportFor(&info->ops[0], "");
ProvideFilesystemSupportFor(&info->ops[1], "file");
}
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h
+31
View File
@@ -0,0 +1,31 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
// Initialize the POSIX filesystem.
//
// In general, the `TF_InitPlugin` symbol doesn't need to be exposed in a header
// file, since the plugin registration will look for the symbol in the DSO file
// that provides the filesystem functionality. However, the POSIX filesystem
// needs to be statically registered in some tests and utilities for building
// the API files at the time of creating the pip package. Hence, we need to
// expose this function so that this filesystem can be statically registered
// when needed.
void TF_InitPlugin(TF_FilesystemPluginInfo* info);
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_H_
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.cc
+65
View File
@@ -0,0 +1,65 @@
/* Copyright 2019 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/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h"
#include <errno.h>
#include <fcntl.h>
#include <limits.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include "tensorflow/c/experimental/filesystem/plugins/posix/copy_file.h"
namespace tf_posix_filesystem {
int TransferFileContents(const char* src, const char* dst, mode_t mode,
off_t size) {
int src_fd = open(src, O_RDONLY);
if (src_fd < 0) return -1;
// When creating file, use the same permissions as original
mode_t open_mode = mode & (S_IRWXU | S_IRWXG | S_IRWXO);
// O_WRONLY | O_CREAT | O_TRUNC:
// Open file for write and if file does not exist, create the file.
// If file exists, truncate its size to 0.
int dst_fd = open(dst, O_WRONLY | O_CREAT | O_TRUNC, open_mode);
if (dst_fd < 0) {
close(src_fd);
return -1;
}
// Both files have been opened, do the transfer.
// Since errno would be overridden by `close` below, save it here.
int error_code = 0;
if (CopyFileContents(dst_fd, src_fd, size) < 0) error_code = errno;
close(src_fd);
close(dst_fd);
if (error_code != 0) {
errno = error_code;
return -1;
} else {
return 0;
}
}
int RemoveSpecialDirectoryEntries(const struct dirent* entry) {
return strcmp(entry->d_name, ".") != 0 && strcmp(entry->d_name, "..") != 0;
}
} // namespace tf_posix_filesystem
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_helper.h
+37
View File
@@ -0,0 +1,37 @@
/* Copyright 2019 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_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
#define TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
#include <dirent.h>
#include <sys/stat.h>
namespace tf_posix_filesystem {
// Copies up to `size` of `src` to `dst`, creating destination if needed.
//
// Callers should pass size of `src` in `size` and the permissions of `src` in
// `mode`. The later is only used if `dst` needs to be created.
int TransferFileContents(const char* src, const char* dst, mode_t mode,
off_t size);
// Returns true only if `entry` points to an entry other than `.` or `..`.
//
// This is a filter for `scandir`.
int RemoveSpecialDirectoryEntries(const struct dirent* entry);
} // namespace tf_posix_filesystem
#endif // TENSORFLOW_C_EXPERIMENTAL_FILESYSTEM_PLUGINS_POSIX_POSIX_FILESYSTEM_HELPER_H_
tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem_static.cc
+41
View File
@@ -0,0 +1,41 @@
/* 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 "absl/log/log.h"
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
#include "tensorflow/c/experimental/filesystem/modular_filesystem_registration.h"
#include "tensorflow/c/experimental/filesystem/plugins/posix/posix_filesystem.h"
#include "tensorflow/core/platform/status.h"
namespace tensorflow {
// Register the POSIX filesystems statically.
// Return value will be unused
bool StaticallyRegisterLocalFilesystems() {
TF_FilesystemPluginInfo info;
TF_InitPlugin(&info);
absl::Status status =
filesystem_registration::RegisterFilesystemPluginImpl(&info);
if (!status.ok()) {
VLOG(0) << "Static POSIX filesystem could not be registered: " << status;
return false;
}
return true;
}
// Perform the actual registration
static bool unused = StaticallyRegisterLocalFilesystems();
} // namespace tensorflow
tensorflow/c/experimental/filesystem/plugins/windows/BUILD
+35
View File
@@ -0,0 +1,35 @@
# Experimental windows filesystem plugin.
load("//tensorflow:tensorflow.bzl", "get_win_copts", "tf_cc_shared_object")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
# Filesystem implementation for Windows environment
tf_cc_shared_object(
name = "windows_filesystem.dll",
framework_so = [],
linkstatic = False,
tags = [
"manual",
"nobuilder",
"notap",
],
visibility = ["//visibility:public"],
deps = [":windows_filesystem_impl"],
)
# The real implementation of the filesystem.
cc_library(
name = "windows_filesystem_impl",
srcs = ["windows_filesystem.cc"],
copts = get_win_copts(),
tags = [
"manual",
"nobuilder",
"notap",
],
deps = ["//tensorflow/c/experimental/filesystem:filesystem_interface"],
)
tensorflow/c/experimental/filesystem/plugins/windows/windows_filesystem.cc
+72
View File
@@ -0,0 +1,72 @@
/* 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 <stdlib.h>
#include <string.h>
#include "tensorflow/c/experimental/filesystem/filesystem_interface.h"
// Implementation of a filesystem for POSIX environments.
// This filesystem will support `file://` and empty (local) URI schemes.
static void* plugin_memory_allocate(size_t size) { return calloc(1, size); }
static void plugin_memory_free(void* ptr) { free(ptr); }
// SECTION 1. Implementation for `TF_RandomAccessFile`
// ----------------------------------------------------------------------------
namespace tf_random_access_file {
// TODO(b/139060984): Implement later
} // namespace tf_random_access_file
// SECTION 2. Implementation for `TF_WritableFile`
// ----------------------------------------------------------------------------
namespace tf_writable_file {
// TODO(b/139060984): Implement later
} // namespace tf_writable_file
// SECTION 3. Implementation for `TF_ReadOnlyMemoryRegion`
// ----------------------------------------------------------------------------
namespace tf_read_only_memory_region {
// TODO(b/139060984): Implement later
} // namespace tf_read_only_memory_region
// SECTION 4. Implementation for `TF_Filesystem`, the actual filesystem
// ----------------------------------------------------------------------------
namespace tf_windows_filesystem {
// TODO(b/139060984): Implement later
} // namespace tf_windows_filesystem
static void ProvideFilesystemSupportFor(TF_FilesystemPluginOps* ops,
const char* uri) {
TF_SetFilesystemVersionMetadata(ops);
ops->scheme = strdup(uri);
}
void TF_InitPlugin(TF_FilesystemPluginInfo* info) {
info->plugin_memory_allocate = plugin_memory_allocate;
info->plugin_memory_free = plugin_memory_free;
info->num_schemes = 2;
info->ops = static_cast<TF_FilesystemPluginOps*>(
plugin_memory_allocate(info->num_schemes * sizeof(info->ops[0])));
ProvideFilesystemSupportFor(&info->ops[0], "");
ProvideFilesystemSupportFor(&info->ops[1], "file");
}
tensorflow/c/experimental/gradients/BUILD
+254
View File
@@ -0,0 +1,254 @@
load(
"//tensorflow:tensorflow.bzl",
"if_libtpu",
"tf_cuda_cc_test",
)
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load(
"//tensorflow/core/platform:build_config_root.bzl",
"tf_cuda_tests_tags",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
# Library of gradient functions.
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "array_grad",
srcs = ["array_grad.cc"],
hdrs = [
"array_grad.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:gradients_internal",
],
)
cc_library(
name = "math_grad",
srcs = ["math_grad.cc"],
hdrs = [
"math_grad.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/c/experimental/ops:array_ops",
"//tensorflow/c/experimental/ops:math_ops",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core/common_runtime/eager:attr_builder",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/types:span",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "nn_grad",
srcs = ["nn_grad.cc"],
hdrs = [
"nn_grad.h",
],
visibility = [
"//tensorflow/python/framework/experimental:__pkg__",
],
deps = [
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/c/eager:immediate_execution_context",
"//tensorflow/c/eager:immediate_execution_tensor_handle",
"//tensorflow/c/experimental/ops:array_ops",
"//tensorflow/c/experimental/ops:math_ops",
"//tensorflow/c/experimental/ops:nn_ops",
"//tensorflow/core/lib/llvm_rtti",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "not_differentiable",
srcs = ["not_differentiable.cc"],
hdrs = [
"not_differentiable.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:gradients_internal",
],
)
cc_library(
name = "gradients",
hdrs = [
"array_grad.h",
"math_grad.h",
"nn_grad.h",
"not_differentiable.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
":array_grad",
":math_grad",
":nn_grad",
":not_differentiable",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:gradients_internal",
],
)
tf_cuda_cc_test(
name = "custom_gradient_test",
size = "small",
srcs = [
"custom_gradient_test.cc",
],
args = ["--heap_check="], # TODO(b/174752220): Remove
tags = tf_cuda_tests_tags(),
deps = [
"//tensorflow/c:tf_status_helper",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:c_api",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/c/eager:c_api_unified_internal",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/c/eager:unified_api_testutil",
"//tensorflow/c/experimental/ops",
"//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:errors",
],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"array_grad.h",
"math_grad.h",
"nn_grad.h",
"not_differentiable.h",
],
visibility = ["//tensorflow/python:__pkg__"],
)
cc_library(
name = "grad_test_helper",
testonly = True,
srcs = ["grad_test_helper.cc"],
hdrs = ["grad_test_helper.h"],
visibility = ["//visibility:private"],
deps = [
"//tensorflow/c/eager:gradient_checker",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/c/eager:unified_api_testutil",
"//tensorflow/c/experimental/gradients/tape:tape_context",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
],
)
tf_cuda_cc_test(
name = "nn_grad_test",
size = "small",
srcs = [
"nn_grad_test.cc",
],
args = ["--heap_check="], # TODO(b/174752220): Remove
tags = tf_cuda_tests_tags() + ["no_cuda_asan"], # b/173654156,
deps = [
":grad_test_helper",
":nn_grad",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c/eager:c_api_test_util",
"//tensorflow/c/eager:c_api_unified_internal",
"//tensorflow/c/eager:unified_api_testutil",
"//tensorflow/c/experimental/gradients/tape:tape_context",
"//tensorflow/c/experimental/ops:nn_ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:tensor_float_32_utils",
] + if_libtpu(
if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
if_true = [],
),
)
tf_cuda_cc_test(
name = "math_grad_test",
size = "small",
srcs = [
"math_grad_test.cc",
],
args = ["--heap_check="], # TODO(b/174752220): Remove
tags = tf_cuda_tests_tags() + ["no_cuda_asan"], # b/173654156,
deps = [
":grad_test_helper",
":math_grad",
"//tensorflow/c:tf_datatype_hdrs",
"//tensorflow/c:tf_status",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:c_api_experimental",
"//tensorflow/c/eager:c_api_test_util",
"//tensorflow/c/eager:c_api_unified_internal",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/c/eager:unified_api_testutil",
"//tensorflow/c/experimental/gradients/tape:tape_context",
"//tensorflow/c/experimental/ops:math_ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:tensor_float_32_utils",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
] + if_libtpu(
if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
if_true = [],
),
)
tf_cuda_cc_test(
name = "array_grad_test",
size = "small",
srcs = [
"array_grad_test.cc",
],
args = ["--heap_check="], # TODO(b/174752220): Remove
tags = tf_cuda_tests_tags() + ["no_cuda_asan"], # b/173654156,
deps = [
":array_grad",
":grad_test_helper",
"//tensorflow/c:tf_status_helper",
"//tensorflow/c/eager:c_api_test_util",
"//tensorflow/c/eager:c_api_unified_internal",
"//tensorflow/c/eager:unified_api_testutil",
"//tensorflow/c/experimental/gradients/tape:tape_context",
"//tensorflow/c/experimental/ops:array_ops",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/platform:tensor_float_32_utils",
] + if_libtpu(
if_false = ["//tensorflow/compiler/mlir/tensorflow/c:mlir_c_api_registration"],
if_true = [],
),
)
tensorflow/c/experimental/gradients/array_grad.cc
+47
View File
@@ -0,0 +1,47 @@
/* 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/c/experimental/gradients/array_grad.h"
#include "tensorflow/c/eager/abstract_context.h"
namespace tensorflow {
namespace gradients {
namespace {
class IdentityNGradientFunction : public GradientFunction {
public:
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
for (int i = 0; i < grad_outputs.size(); i++) {
auto grad_input = grad_outputs[i];
// TODO(srbs): Should we add a copy contructor to AbstractTensorHandle
// that takes care of this similar to `Tensor`?
if (grad_input) {
grad_input->Ref();
}
grad_inputs[i] = grad_input;
}
return absl::OkStatus();
}
~IdentityNGradientFunction() override {}
};
} // namespace
GradientFunction* IdentityNRegisterer(const ForwardOperation& op) {
return new IdentityNGradientFunction;
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/array_grad.h
+26
View File
@@ -0,0 +1,26 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_
#include "tensorflow/c/eager/gradients.h"
namespace tensorflow {
namespace gradients {
GradientFunction* IdentityNRegisterer(const ForwardOperation& op);
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_ARRAY_GRAD_H_
tensorflow/c/experimental/gradients/array_grad_test.cc
+133
View File
@@ -0,0 +1,133 @@
/* 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/c/experimental/gradients/array_grad.h"
#include "tensorflow/c/eager/c_api_test_util.h"
#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
#include "tensorflow/c/eager/unified_api_testutil.h"
#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/platform/tensor_float_32_utils.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace gradients {
namespace internal {
namespace {
using tensorflow::TF_StatusPtr;
absl::Status IdentityNModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
std::vector<AbstractTensorHandle*> temp_outputs(2);
TF_RETURN_IF_ERROR(
ops::IdentityN(ctx, inputs, absl::MakeSpan(temp_outputs), "IdentityN"));
// Although, `ops::IdentityN` returns 2 tensors, the first tensor isn't needed
// for computing gradient so we could safely drop it.
outputs[0] = temp_outputs[1];
temp_outputs[0]->Unref();
return absl::OkStatus();
}
class CppGradients
: public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
protected:
void SetUp() override {
TF_StatusPtr status(TF_NewStatus());
TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
status_ = StatusFromTF_Status(status.get());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
{
AbstractContext* ctx_raw = nullptr;
status_ =
BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
immediate_execution_ctx_.reset(ctx_raw);
}
// Computing numerical gradients with TensorFloat-32 is numerically
// unstable. Some forward pass tests also fail with TensorFloat-32 due to
// low tolerances
enable_tensor_float_32_execution(false);
}
AbstractContextPtr immediate_execution_ctx_;
GradientRegistry registry_;
absl::Status status_;
public:
bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
bool UseFunction() const { return std::get<2>(GetParam()); }
};
TEST_P(CppGradients, TestIdentityNGrad) {
// This test is interesting because the current implementation of GradientTape
// would return [0, 1] whereas we use build_default_zeros_grads=false here
// so we get back [nullptr, 1].
AbstractTensorHandlePtr x1;
{
AbstractTensorHandle* x1_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 1.0f, &x1_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x1.reset(x1_raw);
}
AbstractTensorHandlePtr x2;
{
AbstractTensorHandle* x2_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 1.0f, &x2_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x2.reset(x2_raw);
}
status_ = registry_.Register("IdentityN", IdentityNRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
auto IdentityNGradModel = BuildGradModel(IdentityNModel, registry_);
std::vector<AbstractTensorHandle*> outputs(2);
status_ =
RunModel(IdentityNGradModel, immediate_execution_ctx_.get(),
{x1.get(), x2.get()}, absl::MakeSpan(outputs), UseFunction());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
EXPECT_EQ(outputs[0], nullptr);
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], {1.0f}, /*dims*/ {},
/*abs_error*/ 0));
outputs[1]->Unref();
}
#ifdef PLATFORM_GOOGLE
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#endif
} // namespace
} // namespace internal
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/custom_gradient_test.cc
+138
View File
@@ -0,0 +1,138 @@
/* 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 <memory>
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/c_api.h"
#include "tensorflow/c/eager/c_api_unified_experimental.h"
#include "tensorflow/c/eager/c_api_unified_experimental_internal.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/eager/unified_api_testutil.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace gradients {
namespace internal {
namespace {
using std::vector;
class CustomGradientTest
: public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
protected:
void SetUp() override {
TF_StatusPtr status(TF_NewStatus());
TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
absl::Status s = StatusFromTF_Status(status.get());
CHECK_EQ(errors::OK, s.code()) << s.message();
}
};
class PassThroughGradientFunction : public GradientFunction {
public:
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
CHECK_EQ(grad_outputs.size(), 1);
CHECK_EQ(grad_inputs.size(), 1);
grad_inputs[0] = grad_outputs[0];
if (grad_inputs[0]) {
grad_inputs[0]->Ref();
}
return absl::OkStatus();
}
};
// Computes:
//
// @tf.custom_gradient
// def f(input):
// def grad(grads):
// return grads[0]
// return tf.exp(input), grad
// outputs = [f(inputs[0])]
absl::Status ExpWithPassThroughGrad(
AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
Tape tape(/*persistent=*/false);
tape.Watch(inputs[0]); // Watch x.
AbstractTensorHandle* exp_output;
TF_RETURN_IF_ERROR(ops::Exp(ctx, inputs[0], &exp_output, "Exp"));
std::unique_ptr<GradientFunction> gradient_function(
new PassThroughGradientFunction);
tape.RecordOperation(inputs, {exp_output}, gradient_function.release());
TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx,
/*targets*/ {exp_output},
/*sources=*/inputs,
/*output_gradients=*/{},
/*result=*/outputs));
exp_output->Unref();
return absl::OkStatus();
}
TEST_P(CustomGradientTest, ExpWithPassThroughGrad) {
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
TF_NewStatus(), TF_DeleteStatus);
AbstractContextPtr ctx;
{
AbstractContext* ctx_raw = nullptr;
absl::Status s =
BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
ASSERT_EQ(errors::OK, s.code()) << s.message();
ctx.reset(ctx_raw);
}
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
absl::Status s =
TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 1.0f, &x_raw);
ASSERT_EQ(errors::OK, s.code()) << s.message();
x.reset(x_raw);
}
// Pseudo-code:
//
// tape.watch(x)
// y = exp(x)
// outputs = tape.gradient(y, x)
std::vector<AbstractTensorHandle*> outputs(1);
absl::Status s = RunModel(ExpWithPassThroughGrad, ctx.get(), {x.get()},
absl::MakeSpan(outputs),
/*use_function=*/!std::get<2>(GetParam()));
ASSERT_EQ(errors::OK, s.code()) << s.message();
TF_Tensor* result_tensor;
s = GetValue(outputs[0], &result_tensor);
ASSERT_EQ(errors::OK, s.code()) << s.message();
auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
EXPECT_EQ(*result_value, 1.0);
outputs[0]->Unref();
TF_DeleteTensor(result_tensor);
result_tensor = nullptr;
}
INSTANTIATE_TEST_SUITE_P(
CustomGradientTest, CustomGradientTest,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*executing_eagerly*/ ::testing::Values(true, false)));
} // namespace
} // namespace internal
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/grad_test_helper.cc
+136
View File
@@ -0,0 +1,136 @@
/* 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/c/experimental/gradients/grad_test_helper.h"
#include "tensorflow/c/eager/gradient_checker.h"
#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace gradients {
namespace internal {
void CompareNumericalAndAutodiffGradients(
Model model, Model grad_model, AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs, bool use_function,
double abs_error) {
auto num_inputs = inputs.size();
std::vector<AbstractTensorHandle*> outputs(num_inputs);
auto s = RunModel(grad_model, ctx, inputs, absl::MakeSpan(outputs),
/*use_function=*/use_function);
ASSERT_EQ(errors::OK, s.code()) << s.message();
for (int i = 0; i < num_inputs; ++i) {
if (!outputs[i]) continue;
AbstractTensorHandlePtr numerical_grad;
{
AbstractTensorHandle* numerical_grad_raw;
s = CalcNumericalGrad(ctx, model, inputs,
/*input_index=*/i, use_function,
&numerical_grad_raw);
ASSERT_EQ(errors::OK, s.code()) << s.message();
numerical_grad.reset(numerical_grad_raw);
}
TF_Tensor* numerical_tensor;
s = GetValue(numerical_grad.get(), &numerical_tensor);
ASSERT_EQ(errors::OK, s.code()) << s.message();
auto num_elem_numerical = TF_TensorElementCount(numerical_tensor);
TF_Tensor* analytical_tensor;
s = GetValue(outputs[i], &analytical_tensor);
ASSERT_EQ(errors::OK, s.code()) << s.message();
auto num_elem_analytical = TF_TensorElementCount(analytical_tensor);
ASSERT_EQ(num_elem_numerical, num_elem_analytical);
float* dnumerical = new float[num_elem_numerical]{0};
memcpy(&dnumerical[0], TF_TensorData(numerical_tensor),
TF_TensorByteSize(numerical_tensor));
float* danalytical = new float[num_elem_analytical]{0};
memcpy(&danalytical[0], TF_TensorData(analytical_tensor),
TF_TensorByteSize(analytical_tensor));
for (int j = 0; j < num_elem_numerical; j++) {
ASSERT_NEAR(dnumerical[j], danalytical[j], abs_error);
}
TF_DeleteTensor(analytical_tensor);
TF_DeleteTensor(numerical_tensor);
delete[] danalytical;
delete[] dnumerical;
outputs[i]->Unref();
}
}
void CheckTensorValue(AbstractTensorHandle* t, absl::Span<const float> manuals,
absl::Span<const int64_t> dims, double abs_error) {
TF_Tensor* analytical_tensor;
auto s = GetValue(t, &analytical_tensor);
ASSERT_EQ(errors::OK, s.code()) << s.message();
int64_t num_elem_analytical = 1;
auto num_dims_analytical = TF_NumDims(analytical_tensor);
ASSERT_EQ(dims.size(), num_dims_analytical);
for (int j = 0; j < num_dims_analytical; j++) {
auto dim_analytical = TF_Dim(analytical_tensor, j);
ASSERT_EQ(dims[j], dim_analytical);
num_elem_analytical *= dim_analytical;
}
float* danalytical = new float[num_elem_analytical]{0};
memcpy(&danalytical[0], TF_TensorData(analytical_tensor),
TF_TensorByteSize(analytical_tensor));
for (int64_t j = 0; j < num_elem_analytical; j++) {
if (abs_error == 0) {
ASSERT_EQ(manuals[j], danalytical[j]);
} else {
ASSERT_NEAR(manuals[j], danalytical[j], abs_error);
}
}
TF_DeleteTensor(analytical_tensor);
delete[] danalytical;
}
Model BuildGradModel(Model forward, GradientRegistry registry) {
return [forward_model = std::move(forward),
grad_registry = std::move(registry)](
AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
Tape tape(/*persistent=*/false);
for (size_t i{}; i < inputs.size(); ++i) {
tape.Watch(inputs[i]);
}
std::vector<AbstractTensorHandle*> temp_outputs(1);
AbstractContextPtr tape_ctx(new TapeContext(ctx, &tape, grad_registry));
TF_RETURN_IF_ERROR(
forward_model(tape_ctx.get(), inputs, absl::MakeSpan(temp_outputs)));
TF_RETURN_IF_ERROR(tape.ComputeGradient(ctx, /*targets=*/temp_outputs,
/*sources=*/inputs,
/*output_gradients=*/{}, outputs));
for (auto temp_output : temp_outputs) {
temp_output->Unref();
}
return absl::OkStatus();
};
}
} // namespace internal
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/grad_test_helper.h
+39
View File
@@ -0,0 +1,39 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/eager/unified_api_testutil.h"
namespace tensorflow {
namespace gradients {
namespace internal {
void CompareNumericalAndAutodiffGradients(
Model model, Model grad_model, AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs, bool use_function,
double abs_error = 1e-2);
void CheckTensorValue(AbstractTensorHandle* t, absl::Span<const float> manuals,
absl::Span<const int64_t> dims, double abs_error = 1e-2);
Model BuildGradModel(Model forward, GradientRegistry registry);
} // namespace internal
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_GRAD_TEST_HELPER_H_
tensorflow/c/experimental/gradients/math_grad.cc
+517
View File
@@ -0,0 +1,517 @@
/* 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/c/experimental/gradients/math_grad.h"
#include <string>
#include <vector>
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/common_runtime/eager/attr_builder.h"
#include "tensorflow/core/framework/types.h"
using std::vector;
using tensorflow::ops::AddV2;
using tensorflow::ops::Div;
using tensorflow::ops::DivNoNan;
using tensorflow::ops::MatMul;
using tensorflow::ops::Mul;
using tensorflow::ops::Neg;
using tensorflow::ops::OnesLike;
using tensorflow::ops::SqrtGrad;
namespace tensorflow {
namespace gradients {
namespace {
static absl::Status SafeConj(AbstractContext* ctx,
AbstractTensorHandle* const input,
AbstractTensorHandle** output, const char* name) {
auto dtype = input->DataType();
if (DataTypeIsFloating(BaseType(dtype)) ||
DataTypeIsInteger(BaseType(dtype))) {
return tensorflow::ops::Identity(ctx, input, output, name);
} else if (!DataTypeIsComplex(BaseType(dtype)) &&
BaseType(dtype) != DT_VARIANT) {
return absl::InvalidArgumentError(
absl::StrCat("Expected numeric or variant tensor, got dtype ", dtype));
}
return tensorflow::ops::Conj(ctx, input, output, name);
}
class AddGradientFunction : public GradientFunction {
public:
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
// TODO(b/161805092): Support broadcasting.
DCHECK(grad_outputs[0]);
grad_inputs[0] = grad_outputs[0];
grad_inputs[1] = grad_outputs[0];
grad_inputs[0]->Ref();
grad_inputs[1]->Ref();
return absl::OkStatus();
}
~AddGradientFunction() override = default;
};
class ExpGradientFunction : public GradientFunction {
public:
explicit ExpGradientFunction(AbstractTensorHandle* exp) : exp_(exp) {
exp->Ref();
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
AbstractTensorHandle* conj_output;
std::string name = "Conj_Exp_Grad";
TF_RETURN_IF_ERROR(SafeConj(ctx, exp_.get(), &conj_output, name.c_str()));
AbstractTensorHandlePtr conj_output_releaser(conj_output);
name = "Mul_Exp_Grad";
TF_RETURN_IF_ERROR(
Mul(ctx, conj_output, grad_outputs[0], &grad_inputs[0], name.c_str()));
return absl::OkStatus();
}
~ExpGradientFunction() override = default;
private:
AbstractTensorHandlePtr exp_;
};
class SqrtGradientFunction : public GradientFunction {
public:
explicit SqrtGradientFunction(AbstractTensorHandle* sqrt) : sqrt_(sqrt) {
sqrt->Ref();
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
std::string name = "Sqrt_Grad";
TF_RETURN_IF_ERROR(SqrtGrad(ctx, sqrt_.get(), grad_outputs[0],
&grad_inputs[0], name.c_str()));
return absl::OkStatus();
}
~SqrtGradientFunction() override = default;
private:
AbstractTensorHandlePtr sqrt_;
};
class MatMulGradientFunction : public GradientFunction {
public:
explicit MatMulGradientFunction(vector<AbstractTensorHandle*> f_inputs,
AttrBuilder f_attrs)
: forward_inputs_(f_inputs), forward_attrs_(f_attrs) {
for (auto input : forward_inputs_) {
if (input) {
input->Ref();
}
}
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
/* Given upstream grad U and a matmul op A*B, the gradients are:
*
* dA = U * B.T
* dB = A.T * U
*
* where A.T means `transpose(A)`
*/
AbstractTensorHandle* upstream_grad = grad_outputs[0];
// Get transpose attrs
bool t_a;
TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_a", &t_a));
bool t_b;
TF_RETURN_IF_ERROR(forward_attrs_.Get("transpose_b", &t_b));
// Conj each input
AbstractTensorHandle* conj_output;
std::string name = "Conj_A_MatMul_Grad";
TF_RETURN_IF_ERROR(
SafeConj(ctx, forward_inputs_[0], &conj_output, name.c_str()));
AbstractTensorHandlePtr A(conj_output);
name = "Conj_B_MatMul_Grad";
TF_RETURN_IF_ERROR(
SafeConj(ctx, forward_inputs_[1], &conj_output, name.c_str()));
AbstractTensorHandlePtr B(conj_output);
// Calc Grad
AbstractTensorHandle* matmul_A_output;
AbstractTensorHandle* matmul_B_output;
std::string name_grad_A = "MatMul_Grad_A";
std::string name_grad_B = "MatMul_Grad_B";
if (!t_a && !t_b) {
TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
/*transpose_a = */ false,
/*transpose_b = */ true,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_A.c_str()));
TF_RETURN_IF_ERROR(MatMul(ctx, A.get(), upstream_grad, &matmul_B_output,
/*transpose_a = */ true,
/*transpose_b = */ false,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_B.c_str()));
} else if (!t_a && t_b) {
TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, B.get(), &matmul_A_output,
/*transpose_a = */ false,
/*transpose_b = */ false,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_A.c_str()));
TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, A.get(), &matmul_B_output,
/*transpose_a = */ true,
/*transpose_b = */ false,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_B.c_str()));
} else if (t_a && !t_b) {
TF_RETURN_IF_ERROR(MatMul(ctx, B.get(), upstream_grad, &matmul_A_output,
/*transpose_a = */ false,
/*transpose_b = */ true,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_A.c_str()));
TF_RETURN_IF_ERROR(MatMul(ctx, A.get(), upstream_grad, &matmul_B_output,
/*transpose_a = */ false,
/*transpose_b = */ false,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_B.c_str()));
} else { // t_a && t_b
TF_RETURN_IF_ERROR(MatMul(ctx, B.get(), upstream_grad, &matmul_A_output,
/*transpose_a = */ true,
/*transpose_b = */ true,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_A.c_str()));
TF_RETURN_IF_ERROR(MatMul(ctx, upstream_grad, A.get(), &matmul_B_output,
/*transpose_a = */ true,
/*transpose_b = */ true,
/*grad_a = */ false, /*grad_b = */ false,
name_grad_B.c_str()));
}
// Gradient for A
grad_inputs[0] = matmul_A_output;
// Gradient for B
grad_inputs[1] = matmul_B_output;
return absl::OkStatus();
}
~MatMulGradientFunction() override {
for (auto input : forward_inputs_) {
if (input) {
input->Unref();
}
}
}
private:
// TODO(b/174778737): Only hold needed inputs.
vector<AbstractTensorHandle*> forward_inputs_;
AttrBuilder forward_attrs_;
};
class NegGradientFunction : public GradientFunction {
public:
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
/* Given upstream grad U and a Neg op Y = -X, the gradients are:
*
* dX = -U
*
*/
std::string name = "Neg_Grad";
TF_RETURN_IF_ERROR(
ops::Neg(ctx, grad_outputs[0], &grad_inputs[0], name.c_str()));
return absl::OkStatus();
}
~NegGradientFunction() override = default;
};
class SubGradientFunction : public GradientFunction {
public:
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
/* Given upstream grad U and a Sub op A-B, the gradients are:
*
* dA = U
* dB = -U
*
*/
// Grad for A
DCHECK(grad_outputs[0]);
grad_inputs[0] = grad_outputs[0];
grad_inputs[0]->Ref();
// Grad for B
// negate the upstream grad
std::string name = "Neg_Sub_Grad_B";
TF_RETURN_IF_ERROR(
ops::Neg(ctx, grad_outputs[0], &grad_inputs[1], name.c_str()));
return absl::OkStatus();
}
~SubGradientFunction() override = default;
};
class MulGradientFunction : public GradientFunction {
public:
explicit MulGradientFunction(vector<AbstractTensorHandle*> f_inputs)
: forward_inputs_(f_inputs) {
for (auto input : forward_inputs_) {
if (input) {
input->Ref();
}
}
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
/* Given upstream grad U and a mul op A*B, the gradients are:
*
* dA = U * B
* dB = A * U
*
*/
AbstractTensorHandle* upstream_grad = grad_outputs[0];
// Gradient for A
std::string name = "Mul_Grad_A";
TF_RETURN_IF_ERROR(Mul(ctx, upstream_grad, forward_inputs_[1],
&grad_inputs[0], name.c_str()));
// Gradient for B
name = "Mul_Grad_B";
TF_RETURN_IF_ERROR(Mul(ctx, forward_inputs_[0], upstream_grad,
&grad_inputs[1], name.c_str()));
return absl::OkStatus();
}
~MulGradientFunction() override {
for (auto input : forward_inputs_) {
if (input) {
input->Unref();
}
}
}
private:
// TODO(b/174778737): Only hold needed inputs.
vector<AbstractTensorHandle*> forward_inputs_;
};
class Log1pGradientFunction : public GradientFunction {
public:
explicit Log1pGradientFunction(vector<AbstractTensorHandle*> f_inputs)
: forward_inputs_(f_inputs) {
for (auto input : forward_inputs_) {
if (input) {
input->Ref();
}
}
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
// TODO(vnvo2409): Add control dependency
/* Given upstream grad U and a Log1p op: Y = log(1 + X), the gradients are:
*
* dX = U / (1 + X)
*
*/
AbstractTensorHandle* upstream_grad = grad_outputs[0];
AbstractTensorHandle* X = forward_inputs_[0];
AbstractTensorHandle* temp_output;
// Calculate conjugate of X
std::string name = "Conj_Log1p_Grad_X";
TF_RETURN_IF_ERROR(SafeConj(ctx, X, &temp_output, name.c_str()));
AbstractTensorHandlePtr Conj_X(temp_output);
// Creates Ones
name = "OnesLike_Log1p_Grad_X";
TF_RETURN_IF_ERROR(OnesLike(ctx, Conj_X.get(), &temp_output, name.c_str()));
AbstractTensorHandlePtr Ones_X(temp_output);
name = "Add_Log1p_Grad_X";
// Calculate 1 + Conj(X)
TF_RETURN_IF_ERROR(
AddV2(ctx, Ones_X.get(), Conj_X.get(), &temp_output, name.c_str()));
AbstractTensorHandlePtr Conj_XP1(temp_output);
name = "Div_Log1p_Grad_X";
// Calculate U / (1 + Conj(X))
TF_RETURN_IF_ERROR(
Div(ctx, upstream_grad, Conj_XP1.get(), &grad_inputs[0], name.c_str()));
return absl::OkStatus();
}
~Log1pGradientFunction() override {
for (auto input : forward_inputs_) {
if (input) {
input->Unref();
}
}
}
private:
// TODO(b/174778737): Only hold needed inputs.
vector<AbstractTensorHandle*> forward_inputs_;
};
class DivNoNanGradientFunction : public GradientFunction {
public:
explicit DivNoNanGradientFunction(vector<AbstractTensorHandle*> f_inputs,
vector<AbstractTensorHandle*> f_outputs)
: forward_inputs_(f_inputs), forward_outputs_(f_outputs) {
for (auto input : forward_inputs_) {
if (input) {
input->Ref();
}
}
for (auto output : forward_outputs_) {
if (output) {
output->Ref();
}
}
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
// TODO(vnvo2409): Add shape broadcasting
/* Given upstream grad U and a Div op: Z = X/Y, the gradients are:
*
* dX = U / Y
* dY = -U*X / Y^2 = (X/Y) * -U / Y = -U*Z / Y
*
*/
AbstractTensorHandle* upstream_grad = grad_outputs[0];
AbstractTensorHandle* Y = forward_inputs_[1];
AbstractTensorHandle* Z = forward_outputs_[0];
// Calculate dX = U / Y
std::string name = "Div_Grad_X";
TF_RETURN_IF_ERROR(
DivNoNan(ctx, upstream_grad, Y, &grad_inputs[0], name.c_str()));
AbstractTensorHandle* temp_output;
// Calculate dY = -U*Z / Y
name = "Neg_Div_Grad_Y";
TF_RETURN_IF_ERROR(Neg(ctx, upstream_grad, &temp_output,
name.c_str())); // -U
AbstractTensorHandlePtr MinusU(temp_output);
name = "Mul_Div_Grad_Y";
TF_RETURN_IF_ERROR(Mul(ctx, MinusU.get(), Z, &temp_output,
name.c_str())); // -U*Z
AbstractTensorHandlePtr UZ(temp_output);
name = "Div_Grad_Y";
TF_RETURN_IF_ERROR(DivNoNan(ctx, UZ.get(), Y, &grad_inputs[1],
name.c_str())); // -U*Z / Y
return absl::OkStatus();
}
~DivNoNanGradientFunction() override {
for (auto input : forward_inputs_) {
if (input) {
input->Unref();
}
}
for (auto output : forward_outputs_) {
if (output) {
output->Unref();
}
}
}
private:
// TODO(b/174778737): Only hold needed inputs and outputs.
vector<AbstractTensorHandle*> forward_inputs_;
vector<AbstractTensorHandle*> forward_outputs_;
};
} // namespace
GradientFunction* AddRegisterer(const ForwardOperation& op) {
return new AddGradientFunction;
}
GradientFunction* ExpRegisterer(const ForwardOperation& op) {
return new ExpGradientFunction(op.outputs[0]);
}
GradientFunction* MatMulRegisterer(const ForwardOperation& op) {
return new MatMulGradientFunction(op.inputs, op.attrs);
}
GradientFunction* SqrtRegisterer(const ForwardOperation& op) {
return new SqrtGradientFunction(op.outputs[0]);
}
GradientFunction* NegRegisterer(const ForwardOperation& op) {
return new NegGradientFunction;
}
GradientFunction* SubRegisterer(const ForwardOperation& op) {
return new SubGradientFunction;
}
GradientFunction* MulRegisterer(const ForwardOperation& op) {
return new MulGradientFunction(op.inputs);
}
GradientFunction* Log1pRegisterer(const ForwardOperation& op) {
return new Log1pGradientFunction(op.inputs);
}
GradientFunction* DivNoNanRegisterer(const ForwardOperation& op) {
return new DivNoNanGradientFunction(op.inputs, op.outputs);
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/math_grad.h
+36
View File
@@ -0,0 +1,36 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_
#include "tensorflow/c/eager/gradients.h"
namespace tensorflow {
namespace gradients {
GradientFunction* AddRegisterer(const ForwardOperation& op);
GradientFunction* ExpRegisterer(const ForwardOperation& op);
GradientFunction* MatMulRegisterer(const ForwardOperation& op);
GradientFunction* SqrtRegisterer(const ForwardOperation& op);
GradientFunction* NegRegisterer(const ForwardOperation& op);
GradientFunction* SubRegisterer(const ForwardOperation& op);
GradientFunction* MulRegisterer(const ForwardOperation& op);
GradientFunction* Log1pRegisterer(const ForwardOperation& op);
GradientFunction* DivNoNanRegisterer(const ForwardOperation& op);
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_MATH_GRAD_H_
tensorflow/c/experimental/gradients/math_grad_test.cc
+461
View File
@@ -0,0 +1,461 @@
/* 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/c/experimental/gradients/math_grad.h"
#include <cstdint>
#include <cstring>
#include <tuple>
#include <vector>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/c_api_unified_experimental.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/eager/unified_api_testutil.h"
#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/tf_datatype.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/tensor_float_32_utils.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace gradients {
namespace internal {
namespace {
using tensorflow::TF_StatusPtr;
absl::Status AddModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::AddV2(ctx, inputs[0], inputs[1], &outputs[0], "Add");
}
absl::Status ExpModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Exp(ctx, inputs[0], &outputs[0], "Exp");
}
absl::Status SqrtModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Sqrt(ctx, inputs[0], &outputs[0], "Sqrt");
}
absl::Status NegModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Neg(ctx, inputs[0], &outputs[0], "Neg");
}
absl::Status SubModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Sub(ctx, inputs[0], inputs[1], &outputs[0], "Sub");
}
absl::Status MulModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Mul(ctx, inputs[0], inputs[1], &outputs[0], "Mul");
}
absl::Status Log1pModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Log1p(ctx, inputs[0], &outputs[0], "Log1p");
}
absl::Status DivNoNanModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::DivNoNan(ctx, inputs[0], inputs[1], &outputs[0], "DivNoNan");
}
class CppGradients
: public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
protected:
void SetUp() override {
TF_StatusPtr status(TF_NewStatus());
TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
status_ = StatusFromTF_Status(status.get());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
{
AbstractContext* ctx_raw = nullptr;
status_ =
BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
immediate_execution_ctx_.reset(ctx_raw);
}
// Computing numerical gradients with TensorFloat-32 is numerically
// unstable. Some forward pass tests also fail with TensorFloat-32 due to
// low tolerances
enable_tensor_float_32_execution(false);
}
AbstractContextPtr immediate_execution_ctx_;
GradientRegistry registry_;
absl::Status status_;
public:
bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
bool UseFunction() const { return std::get<2>(GetParam()); }
};
TEST_P(CppGradients, TestAddGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
AbstractTensorHandlePtr y;
{
AbstractTensorHandle* y_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
y.reset(y_raw);
}
// TODO(srbs): Rename ops::Add to ops::AddV2 and AddRegister to
// AddV2Registerer.
status_ = registry_.Register("AddV2", AddRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
AddModel, BuildGradModel(AddModel, registry_),
immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
}
TEST_P(CppGradients, TestExpGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
status_ = registry_.Register("Exp", ExpRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
ExpModel, BuildGradModel(ExpModel, registry_),
immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}
TEST_P(CppGradients, TestMatMulGrad) {
// TODO(vnvo2409): Figure out why `gradient_checker` does not work very
// well with `MatMul` and remove `TestMatMul*` in
// `mnist_gradients_test` when done.
GTEST_SKIP();
float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
int64_t A_dims[] = {3, 3};
AbstractTensorHandlePtr A;
{
AbstractTensorHandle* A_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
A.reset(A_raw);
}
float B_vals[] = {9.0f, 8.0f, 7.0f, 6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
int64_t B_dims[] = {3, 3};
AbstractTensorHandlePtr B;
{
AbstractTensorHandle* B_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), B_vals, B_dims, 2, &B_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
B.reset(B_raw);
}
status_ = registry_.Register("MatMul", MatMulRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
for (bool transpose_a : {false, true}) {
for (bool transpose_b : {false, true}) {
Model MatMulModel =
[transpose_a, transpose_b](
AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0], transpose_a,
transpose_b, /*grad_a=*/false, /*grad_b=*/false,
"MatMul");
};
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
MatMulModel, BuildGradModel(MatMulModel, registry_),
immediate_execution_ctx_.get(), {A.get(), B.get()}, UseFunction()));
}
}
}
TEST_P(CppGradients, TestMatMulGradManual) {
float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f};
int64_t A_dims[] = {3, 3};
AbstractTensorHandlePtr A;
{
AbstractTensorHandle* A_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
A.reset(A_raw);
}
float B_vals[] = {9.0f, 8.0f, 7.0f, 6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
int64_t B_dims[] = {3, 3};
AbstractTensorHandlePtr B;
{
AbstractTensorHandle* B_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), B_vals, B_dims, 2, &B_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
B.reset(B_raw);
}
status_ = registry_.Register("MatMul", MatMulRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
bool transpose_a_vals[] = {false, false, true, true};
bool transpose_b_vals[] = {false, true, false, true};
float dA_vals[4][9] = {{24, 15, 6, 24, 15, 6, 24, 15, 6},
{18, 15, 12, 18, 15, 12, 18, 15, 12},
{24, 24, 24, 15, 15, 15, 6, 6, 6},
{18, 18, 18, 15, 15, 15, 12, 12, 12}};
float dB_vals[4][9] = {{12, 12, 12, 15, 15, 15, 18, 18, 18},
{12, 15, 18, 12, 15, 18, 12, 15, 18},
{6, 6, 6, 15, 15, 15, 24, 24, 24},
{6, 15, 24, 6, 15, 24, 6, 15, 24}};
for (int i{}; i < 4; ++i) {
bool transpose_a = transpose_a_vals[i];
bool transpose_b = transpose_b_vals[i];
Model MatMulModel =
[transpose_a, transpose_b](
AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) -> absl::Status {
return ops::MatMul(ctx, inputs[0], inputs[1], &outputs[0], transpose_a,
transpose_b, /*grad_a=*/false, /*grad_b=*/false,
"MatMul");
};
Model MatMulGradModel = BuildGradModel(MatMulModel, registry_);
std::vector<AbstractTensorHandle*> outputs(2);
status_ =
RunModel(MatMulGradModel, immediate_execution_ctx_.get(),
{A.get(), B.get()}, absl::MakeSpan(outputs), UseFunction());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], dA_vals[i],
/*dims*/ {3, 3},
/*abs_error*/ 0));
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], dB_vals[i],
/*dims*/ {3, 3},
/*abs_error*/ 0));
outputs[0]->Unref();
outputs[1]->Unref();
}
}
TEST_P(CppGradients, TestSqrtGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
status_ = registry_.Register("Sqrt", SqrtRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
SqrtModel, BuildGradModel(SqrtModel, registry_),
immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}
TEST_P(CppGradients, TestNegGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
status_ = registry_.Register("Neg", NegRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
NegModel, BuildGradModel(NegModel, registry_),
immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}
TEST_P(CppGradients, TestSubGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
AbstractTensorHandlePtr y;
{
AbstractTensorHandle* y_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
y.reset(y_raw);
}
status_ = registry_.Register("Sub", SubRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
SubModel, BuildGradModel(SubModel, registry_),
immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
}
TEST_P(CppGradients, TestMulGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
AbstractTensorHandlePtr y;
{
AbstractTensorHandle* y_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
y.reset(y_raw);
}
status_ = registry_.Register("Mul", MulRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
MulModel, BuildGradModel(MulModel, registry_),
immediate_execution_ctx_.get(), {x.get(), y.get()}, UseFunction()));
}
TEST_P(CppGradients, TestLog1pGrad) {
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
status_ = registry_.Register("Log1p", Log1pRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
Log1pModel, BuildGradModel(Log1pModel, registry_),
immediate_execution_ctx_.get(), {x.get()}, UseFunction()));
}
TEST_P(CppGradients, TestDivNoNanGrad) {
status_ = registry_.Register("DivNoNan", DivNoNanRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
auto DivNoNanGradModel = BuildGradModel(DivNoNanModel, registry_);
AbstractTensorHandlePtr x;
{
AbstractTensorHandle* x_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &x_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
x.reset(x_raw);
}
AbstractTensorHandlePtr y;
{
AbstractTensorHandle* y_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 2.0f, &y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
y.reset(y_raw);
}
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
DivNoNanModel, DivNoNanGradModel, immediate_execution_ctx_.get(),
{x.get(), y.get()}, UseFunction()));
// `DivNoNanGradModel` should return {`0`, `0`} when the denominator is `0`.
AbstractTensorHandlePtr z;
{
AbstractTensorHandle* z_raw = nullptr;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 0.0f, &z_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
z.reset(z_raw);
}
std::vector<AbstractTensorHandle*> outputs(2);
status_ =
RunModel(DivNoNanGradModel, immediate_execution_ctx_.get(),
{x.get(), z.get()}, absl::MakeSpan(outputs), UseFunction());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], {0.0f}, /*dims*/ {},
/*abs_error*/ 0));
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[1], {0.0f}, /*dims*/ {},
/*abs_error*/ 0));
outputs[0]->Unref();
outputs[1]->Unref();
}
#ifdef PLATFORM_GOOGLE
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#endif
} // namespace
} // namespace internal
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/nn_grad.cc
+171
View File
@@ -0,0 +1,171 @@
/* 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/c/experimental/gradients/nn_grad.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_context.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "tensorflow/c/experimental/ops/math_ops.h"
#include "tensorflow/c/experimental/ops/nn_ops.h"
#include "tensorflow/core/lib/llvm_rtti/llvm_rtti.h"
#include "tensorflow/core/platform/errors.h"
using std::vector;
using tensorflow::ops::BiasAddGrad;
using tensorflow::ops::ReluGrad;
namespace tensorflow {
namespace gradients {
namespace {
class ReluGradientFunction : public GradientFunction {
public:
explicit ReluGradientFunction(vector<AbstractTensorHandle*> f_outputs)
: forward_outputs_(f_outputs) {
for (auto output : forward_outputs_) {
if (output) {
output->Ref();
}
}
}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
AbstractTensorHandle* upstream_grad = grad_outputs[0];
AbstractTensorHandle* activations = forward_outputs_[0];
// Calculate Grad
std::string name = "relu_grad";
TF_RETURN_IF_ERROR(ReluGrad(ctx, upstream_grad, activations,
&grad_inputs[0], name.c_str()));
return absl::OkStatus();
}
~ReluGradientFunction() override {
for (auto output : forward_outputs_) {
if (output) {
output->Unref();
}
}
}
private:
// TODO(b/174778737): Only hold needed outputs.
vector<AbstractTensorHandle*> forward_outputs_;
};
absl::Status BroadcastMul(AbstractContext* ctx, AbstractTensorHandle* vec,
AbstractTensorHandle* mat,
absl::Span<AbstractTensorHandle*> outputs) {
if (!isa<ImmediateExecutionContext>(ctx)) {
// TODO(b/168850692): Fix this.
return absl::UnimplementedError(
"BroadcastMul is not supported in tracing mode yet.");
}
auto imm_ctx = dyn_cast<ImmediateExecutionContext>(ctx);
AbstractTensorPtr minus_1(imm_ctx->CreateInt32Scalar(-1));
ImmediateTensorHandlePtr dim(imm_ctx->CreateLocalHandle(minus_1.get()));
AbstractTensorHandle* expand_dims_outputs;
TF_RETURN_IF_ERROR(
ops::ExpandDims(ctx, vec, dim.get(), &expand_dims_outputs, "ExpandDims"));
TF_RETURN_IF_ERROR(
ops::Mul(ctx, expand_dims_outputs, mat, &outputs[0], "Mul"));
expand_dims_outputs->Unref();
return absl::OkStatus();
}
class SparseSoftmaxCrossEntropyWithLogitsGradientFunction
: public GradientFunction {
public:
explicit SparseSoftmaxCrossEntropyWithLogitsGradientFunction(
vector<AbstractTensorHandle*> f_outputs)
: forward_outputs_(f_outputs) {}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
// Grad for Softmax Input
TF_RETURN_IF_ERROR(BroadcastMul(
ctx, grad_outputs[0], forward_outputs_[1],
grad_inputs.subspan(0, 1))); // upstream_grad * local softmax grad
// Grad for labels is null
grad_inputs[1] = nullptr;
return absl::OkStatus();
}
~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override {}
private:
vector<AbstractTensorHandle*> forward_outputs_;
};
// TODO(vnvo2409): Add python test
class BiasAddGradientFunction : public GradientFunction {
public:
explicit BiasAddGradientFunction(AttrBuilder f_attrs)
: forward_attrs_(f_attrs) {}
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override {
/* Given upstream grad U and a BiasAdd: A + bias, the gradients are:
*
* dA = U
* dbias = reduceSum(U, dims = channel_dim)
*/
AbstractTensorHandle* upstream_grad = grad_outputs[0];
DCHECK(upstream_grad);
// Recover data format from forward pass for gradient.
std::string data_format;
TF_RETURN_IF_ERROR(forward_attrs_.Get("data_format", &data_format));
// Grad for A
grad_inputs[0] = upstream_grad;
grad_inputs[0]->Ref();
// Grad for bias
std::string name = "bias_add_grad";
TF_RETURN_IF_ERROR(BiasAddGrad(ctx, upstream_grad, &grad_inputs[1],
data_format.c_str(), name.c_str()));
return absl::OkStatus();
}
~BiasAddGradientFunction() override {}
private:
AttrBuilder forward_attrs_;
};
} // namespace
GradientFunction* ReluRegisterer(const ForwardOperation& op) {
return new ReluGradientFunction(op.outputs);
}
GradientFunction* SparseSoftmaxCrossEntropyWithLogitsRegisterer(
const ForwardOperation& op) {
return new SparseSoftmaxCrossEntropyWithLogitsGradientFunction(op.outputs);
}
GradientFunction* BiasAddRegisterer(const ForwardOperation& op) {
return new BiasAddGradientFunction(op.attrs);
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/nn_grad.h
+29
View File
@@ -0,0 +1,29 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_
#include "tensorflow/c/eager/gradients.h"
namespace tensorflow {
namespace gradients {
GradientFunction* ReluRegisterer(const ForwardOperation& op);
GradientFunction* SparseSoftmaxCrossEntropyWithLogitsRegisterer(
const ForwardOperation& op);
GradientFunction* BiasAddRegisterer(const ForwardOperation& op);
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NN_GRAD_H_
tensorflow/c/experimental/gradients/nn_grad_test.cc
+228
View File
@@ -0,0 +1,228 @@
/* 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/c/experimental/gradients/nn_grad.h"
#include "tensorflow/c/eager/c_api_test_util.h"
#include "tensorflow/c/eager/unified_api_testutil.h"
#include "tensorflow/c/experimental/gradients/grad_test_helper.h"
#include "tensorflow/c/experimental/gradients/tape/tape_context.h"
#include "tensorflow/c/experimental/ops/nn_ops.h"
#include "tensorflow/c/tf_status_helper.h"
#include "tensorflow/core/platform/tensor_float_32_utils.h"
#include "tensorflow/core/platform/test.h"
namespace tensorflow {
namespace gradients {
namespace internal {
namespace {
using tensorflow::TF_StatusPtr;
absl::Status ReluModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::Relu(ctx, inputs[0], &outputs[0], "Relu");
}
absl::Status SparseSoftmaxCrossEntropyWithLogitsModel(
AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
AbstractTensorHandle* loss;
AbstractTensorHandle* backprop;
TF_RETURN_IF_ERROR(ops::SparseSoftmaxCrossEntropyWithLogits(
ctx, inputs[0], inputs[1], &loss, &backprop,
"SparseSoftmaxCrossEntropyWithLogits"));
// `gradient_checker` only works with model that returns only 1 tensor.
// Although, `ops::SparseSoftmaxCrossEntropyWithLogits` returns 2 tensors, the
// second tensor isn't needed for computing gradient so we could safely drop
// it.
outputs[0] = loss;
backprop->Unref();
return absl::OkStatus();
}
absl::Status BiasAddModel(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs) {
return ops::BiasAdd(ctx, inputs[0], inputs[1], &outputs[0], "NHWC",
"BiasAdd");
}
class CppGradients
: public ::testing::TestWithParam<std::tuple<const char*, bool, bool>> {
protected:
void SetUp() override {
TF_StatusPtr status(TF_NewStatus());
TF_SetTracingImplementation(std::get<0>(GetParam()), status.get());
status_ = StatusFromTF_Status(status.get());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
{
AbstractContext* ctx_raw = nullptr;
status_ =
BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
immediate_execution_ctx_.reset(ctx_raw);
}
// Computing numerical gradients with TensorFloat-32 is numerically
// unstable. Some forward pass tests also fail with TensorFloat-32 due to
// low tolerances
enable_tensor_float_32_execution(false);
}
AbstractContextPtr immediate_execution_ctx_;
GradientRegistry registry_;
absl::Status status_;
public:
bool UseMlir() const { return strcmp(std::get<0>(GetParam()), "mlir") == 0; }
bool UseFunction() const { return std::get<2>(GetParam()); }
};
TEST_P(CppGradients, TestReluGrad) {
status_ = registry_.Register("Relu", ReluRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
auto ReluGradModel = BuildGradModel(ReluModel, registry_);
float X_vals[] = {1.0f, 2.0f, 3.0f, -5.0f, -4.0f, -3.0f, 2.0f, 10.0f, -1.0f};
int64_t X_dims[] = {3, 3};
AbstractTensorHandlePtr X;
{
AbstractTensorHandle* X_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
X.reset(X_raw);
}
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
ReluModel, ReluGradModel, immediate_execution_ctx_.get(), {X.get()},
UseFunction()));
// Mathematically, Relu isn't differentiable at `0`. So `gradient_checker`
// does not work with it.
AbstractTensorHandlePtr Y;
{
AbstractTensorHandle* Y_raw;
status_ = TestScalarTensorHandle<float, TF_FLOAT>(
immediate_execution_ctx_.get(), 0.0f, &Y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
Y.reset(Y_raw);
}
std::vector<AbstractTensorHandle*> outputs(1);
status_ = RunModel(ReluGradModel, immediate_execution_ctx_.get(), {Y.get()},
absl::MakeSpan(outputs), UseFunction());
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CheckTensorValue(outputs[0], {0.0f}, /*dims*/ {},
/*abs_error*/ 0));
outputs[0]->Unref();
}
TEST_P(CppGradients, TestSparseSoftmaxCrossEntropyWithLogitsGrad) {
if (UseFunction()) {
// TODO(b/168850692): Enable this.
GTEST_SKIP() << "Can't take gradient of "
"SparseSoftmaxCrossEntropyWithLogits in tracing mode.";
}
// Score
float X_vals[] = {1.0f, 2.0f, 3.0f, -5.0f, -4.0f, -3.0f, 2.0f, 0.0f, -1.0f};
int64_t X_dims[] = {3, 3};
AbstractTensorHandlePtr X;
{
AbstractTensorHandle* X_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), X_vals, X_dims, 2, &X_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
X.reset(X_raw);
}
// Label
int32_t Y_vals[] = {1, 0, 1};
int64_t Y_dims[] = {3};
AbstractTensorHandlePtr Y;
{
AbstractTensorHandle* Y_raw;
status_ = TestTensorHandleWithDims<int32_t, TF_INT32>(
immediate_execution_ctx_.get(), Y_vals, Y_dims, 1, &Y_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
Y.reset(Y_raw);
}
status_ = registry_.Register("SparseSoftmaxCrossEntropyWithLogits",
SparseSoftmaxCrossEntropyWithLogitsRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
SparseSoftmaxCrossEntropyWithLogitsModel,
BuildGradModel(SparseSoftmaxCrossEntropyWithLogitsModel, registry_),
immediate_execution_ctx_.get(), {X.get(), Y.get()}, UseFunction()));
}
TEST_P(CppGradients, TestBiasAddGrad) {
if (UseFunction() && UseMlir()) {
GTEST_SKIP() << "SetAttrString has not been implemented yet.\n";
}
// A
float A_vals[] = {1.0f, 2.0f, 3.0f, 4.0f};
int64_t A_dims[] = {2, 2};
AbstractTensorHandlePtr A;
{
AbstractTensorHandle* A_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), A_vals, A_dims, 2, &A_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
A.reset(A_raw);
}
// Bias
float Bias_vals[] = {2.0f, 3.0f};
int64_t Bias_dims[] = {2};
AbstractTensorHandlePtr Bias;
{
AbstractTensorHandle* Bias_raw;
status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
immediate_execution_ctx_.get(), Bias_vals, Bias_dims, 1, &Bias_raw);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
Bias.reset(Bias_raw);
}
status_ = registry_.Register("BiasAdd", BiasAddRegisterer);
ASSERT_EQ(errors::OK, status_.code()) << status_.message();
ASSERT_NO_FATAL_FAILURE(CompareNumericalAndAutodiffGradients(
BiasAddModel, BuildGradModel(BiasAddModel, registry_),
immediate_execution_ctx_.get(), {A.get(), Bias.get()}, UseFunction()));
}
#ifdef PLATFORM_GOOGLE
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#else
INSTANTIATE_TEST_SUITE_P(
UnifiedCAPI, CppGradients,
::testing::Combine(::testing::Values("graphdef", "mlir"),
/*tfrt*/ ::testing::Values(false),
/*use_function*/ ::testing::Values(true, false)));
#endif
} // namespace
} // namespace internal
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/not_differentiable.cc
+35
View File
@@ -0,0 +1,35 @@
/* 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/c/experimental/gradients/not_differentiable.h"
namespace tensorflow {
namespace gradients {
absl::Status NotDifferentiableGradientFunction::Compute(
AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) {
for (int i = 0; i < grad_inputs.size(); i++) {
grad_inputs[i] = nullptr;
}
return absl::OkStatus();
}
absl::Status RegisterNotDifferentiable(GradientRegistry* registry,
const std::string& op) {
return registry->Register(op, [](const ForwardOperation& op) {
return new NotDifferentiableGradientFunction;
});
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/not_differentiable.h
+35
View File
@@ -0,0 +1,35 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/gradients.h"
namespace tensorflow {
namespace gradients {
// Ignores `grad_outputs` and sets all entries in grad_inputs to nullptr.
class NotDifferentiableGradientFunction : public GradientFunction {
absl::Status Compute(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> grad_outputs,
absl::Span<AbstractTensorHandle*> grad_inputs) override;
};
// Shorthand for registry->Register(op, new NotDifferentiableGradientFunction)
absl::Status RegisterNotDifferentiable(GradientRegistry* registry,
const std::string& op);
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_NOT_DIFFERENTIABLE_H_
tensorflow/c/experimental/gradients/tape/BUILD
+93
View File
@@ -0,0 +1,93 @@
# A tape built on top of unified execution APIs.
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "tape_context",
srcs = ["tape_context.cc"],
hdrs = [
"tape_context.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
":tape_operation",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_function",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/status",
],
)
cc_library(
name = "tape_operation",
srcs = ["tape_operation.cc"],
hdrs = [
"tape_operation.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"//tensorflow/core/platform:status",
"//tensorflow/core/platform:stringpiece",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@xla//xla/tsl/platform:errors",
],
)
cc_library(
name = "tape",
hdrs = [
"tape_context.h",
"tape_operation.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
":tape_context",
":tape_operation",
"//tensorflow/c:tensor_interface",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_function",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:gradients_internal",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"tape_context.h",
"tape_operation.h",
],
visibility = [
"//tensorflow:internal",
],
)
tensorflow/c/experimental/gradients/tape/tape_context.cc
+52
View File
@@ -0,0 +1,52 @@
/* 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/c/experimental/gradients/tape/tape_context.h"
#include "absl/status/status.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_function.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/experimental/gradients/tape/tape_operation.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace gradients {
TapeContext::TapeContext(AbstractContext* c, Tape* tape,
const GradientRegistry& registry)
: AbstractContext(kTape), parent_ctx_(c), tape_(tape), registry_(registry) {
// TODO(srbs): Make AbstractContext ref counted.
// parent_ctx_->Ref();
}
void TapeContext::Release() {
// TODO(srbs): Change to Unref()
delete this;
}
TapeContext::~TapeContext() {
// TODO(srbs): Make AbstractContext ref counted.
// parent_ctx_->Unref();
}
TapeOperation* TapeContext::CreateOperation() {
return new TapeOperation(parent_ctx_->CreateOperation(), tape_, registry_);
}
absl::Status TapeContext::RegisterFunction(AbstractFunction* f) {
return parent_ctx_->RegisterFunction(f);
}
absl::Status TapeContext::RemoveFunction(const std::string& func) {
return parent_ctx_->RemoveFunction(func);
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/tape/tape_context.h
+49
View File
@@ -0,0 +1,49 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_
#include "absl/status/status.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_function.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/experimental/gradients/tape/tape_operation.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace gradients {
class TapeContext : public AbstractContext {
public:
explicit TapeContext(AbstractContext*, Tape*, const GradientRegistry&);
void Release() override;
TapeOperation* CreateOperation() override;
absl::Status RegisterFunction(AbstractFunction*) override;
absl::Status RemoveFunction(const std::string& func) override;
// For LLVM style RTTI.
static bool classof(const AbstractContext* ptr) {
return ptr->getKind() == kTape;
}
~TapeContext() override;
private:
AbstractContext* parent_ctx_; // Not owned.
Tape* tape_;
const GradientRegistry& registry_;
};
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_CONTEXT_H_
tensorflow/c/experimental/gradients/tape/tape_operation.cc
+246
View File
@@ -0,0 +1,246 @@
/* 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/c/experimental/gradients/tape/tape_operation.h"
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <vector>
#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_operation.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/tensor_interface.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/tensor_shape.h"
#include "tensorflow/core/framework/tensor_shape.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/lib/gtl/array_slice.h"
#include "tensorflow/core/platform/errors.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/strcat.h"
#include "tensorflow/core/platform/stringpiece.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace gradients {
TapeOperation::TapeOperation(AbstractOperation* parent_op, Tape* tape,
const GradientRegistry& registry)
: AbstractOperation(kTape),
parent_op_(parent_op),
tape_(tape),
registry_(registry) {
// TODO(b/172003047): Consider making AbstractOperation RefCounted.
// parent_op_->Ref();
}
void TapeOperation::Release() {
// TODO(srbs): Change to Unref().
delete this;
}
TapeOperation::~TapeOperation() {
// TODO(b/172003047): Consider making AbstractOperation RefCounted.
// parent_op->Unref();
}
absl::Status TapeOperation::Reset(const char* op, const char* raw_device_name) {
forward_op_.op_name = op;
forward_op_.attrs.Reset(op);
forward_op_.inputs.clear();
forward_op_.outputs.clear();
return parent_op_->Reset(op, raw_device_name);
}
const std::string& TapeOperation::Name() const { return parent_op_->Name(); }
const std::string& TapeOperation::DeviceName() const {
return parent_op_->DeviceName();
}
absl::Status TapeOperation::SetDeviceName(const char* name) {
return parent_op_->SetDeviceName(name);
}
absl::Status TapeOperation::AddInput(AbstractTensorHandle* input) {
TF_RETURN_IF_ERROR(parent_op_->AddInput(input));
forward_op_.inputs.push_back(input);
return absl::OkStatus();
}
absl::Status TapeOperation::AddInputList(
absl::Span<AbstractTensorHandle* const> inputs) {
TF_RETURN_IF_ERROR(parent_op_->AddInputList(inputs));
for (auto input : inputs) {
forward_op_.inputs.push_back(input);
}
return absl::OkStatus();
}
absl::Status TapeOperation::SetAttrString(const char* attr_name,
const char* data, size_t length) {
forward_op_.attrs.Set(attr_name, absl::string_view(data, length));
return parent_op_->SetAttrString(attr_name, data, length);
}
absl::Status TapeOperation::SetAttrInt(const char* attr_name, int64_t value) {
forward_op_.attrs.Set(attr_name, static_cast<int64_t>(value));
return parent_op_->SetAttrInt(attr_name, value);
}
absl::Status TapeOperation::SetAttrFloat(const char* attr_name, float value) {
forward_op_.attrs.Set(attr_name, value);
return parent_op_->SetAttrFloat(attr_name, value);
}
absl::Status TapeOperation::SetAttrBool(const char* attr_name, bool value) {
forward_op_.attrs.Set(attr_name, value);
return parent_op_->SetAttrBool(attr_name, value);
}
absl::Status TapeOperation::SetAttrType(const char* attr_name, DataType value) {
forward_op_.attrs.Set(attr_name, value);
return parent_op_->SetAttrType(attr_name, value);
}
absl::Status TapeOperation::SetAttrShape(const char* attr_name,
const int64_t* dims,
const int num_dims) {
if (num_dims > TensorShape::MaxDimensions()) {
return errors::InvalidArgument("Value specified for `", attr_name, "` has ",
num_dims,
" dimensions which is over the limit of ",
TensorShape::MaxDimensions(), ".");
}
TensorShapeProto proto;
if (num_dims < 0) {
proto.set_unknown_rank(true);
} else {
for (int d = 0; d < num_dims; ++d) {
proto.add_dim()->set_size(dims[d]);
}
}
forward_op_.attrs.Set(attr_name, proto);
return parent_op_->SetAttrShape(attr_name, dims, num_dims);
}
absl::Status TapeOperation::SetAttrFunction(const char* attr_name,
const AbstractOperation* value) {
return tensorflow::errors::Unimplemented(
"SetAttrFunction has not been implemented yet.");
}
absl::Status TapeOperation::SetAttrFunctionName(const char* attr_name,
const char* value,
size_t length) {
return tensorflow::errors::Unimplemented(
"SetAttrFunctionName has not been implemented "
"yet.");
}
absl::Status TapeOperation::SetAttrTensor(const char* attr_name,
AbstractTensorInterface* tensor) {
return tensorflow::errors::Unimplemented(
"SetAttrTensor has not been implemented yet.");
}
absl::Status TapeOperation::SetAttrStringList(const char* attr_name,
const void* const* values,
const size_t* lengths,
int num_values) {
std::vector<absl::string_view> v(num_values);
for (int i = 0; i < num_values; ++i) {
v[i] = absl::string_view(static_cast<const char*>(values[i]), lengths[i]);
}
forward_op_.attrs.Set(attr_name, v);
return parent_op_->SetAttrStringList(attr_name, values, lengths, num_values);
}
absl::Status TapeOperation::SetAttrFloatList(const char* attr_name,
const float* values,
int num_values) {
forward_op_.attrs.Set(attr_name,
gtl::ArraySlice<const float>(values, num_values));
return parent_op_->SetAttrFloatList(attr_name, values, num_values);
}
absl::Status TapeOperation::SetAttrIntList(const char* attr_name,
const int64_t* values,
int num_values) {
forward_op_.attrs.Set(
attr_name, gtl::ArraySlice<const int64_t>(
reinterpret_cast<const int64_t*>(values), num_values));
return parent_op_->SetAttrIntList(attr_name, values, num_values);
}
absl::Status TapeOperation::SetAttrTypeList(const char* attr_name,
const DataType* values,
int num_values) {
forward_op_.attrs.Set(attr_name,
gtl::ArraySlice<const DataType>(values, num_values));
return parent_op_->SetAttrTypeList(attr_name, values, num_values);
}
absl::Status TapeOperation::SetAttrBoolList(const char* attr_name,
const unsigned char* values,
int num_values) {
std::unique_ptr<bool[]> b(new bool[num_values]);
for (int i = 0; i < num_values; ++i) {
b[i] = values[i];
}
forward_op_.attrs.Set(attr_name,
gtl::ArraySlice<const bool>(b.get(), num_values));
return parent_op_->SetAttrBoolList(attr_name, values, num_values);
}
absl::Status TapeOperation::SetAttrShapeList(const char* attr_name,
const int64_t** dims,
const int* num_dims,
int num_values) {
std::unique_ptr<TensorShapeProto[]> proto(new TensorShapeProto[num_values]);
for (int i = 0; i < num_values; ++i) {
const auto num_dims_i = num_dims[i];
if (num_dims_i > TensorShape::MaxDimensions()) {
return errors::InvalidArgument(
strings::StrCat("Value specified for `", attr_name, "` has ",
num_dims_i, " dimensions which is over the limit of ",
TensorShape::MaxDimensions(), "."));
}
if (num_dims_i < 0) {
proto[i].set_unknown_rank(true);
} else {
const int64_t* dims_i = dims[i];
auto proto_i = &proto[i];
for (int d = 0; d < num_dims_i; ++d) {
proto_i->add_dim()->set_size(dims_i[d]);
}
}
}
forward_op_.attrs.Set(
attr_name, absl::Span<const TensorShapeProto>(proto.get(), num_values));
return parent_op_->SetAttrShapeList(attr_name, dims, num_dims, num_values);
}
absl::Status TapeOperation::SetAttrFunctionList(
const char* attr_name, absl::Span<const AbstractOperation*> values) {
return tensorflow::errors::Unimplemented(
"SetAttrFunctionList has not been "
"implemented yet.");
}
AbstractOperation* TapeOperation::GetBackingOperation() { return parent_op_; }
absl::Status TapeOperation::Execute(absl::Span<AbstractTensorHandle*> retvals,
int* num_retvals) {
TF_RETURN_IF_ERROR(parent_op_->Execute(retvals, num_retvals));
for (int i = 0; i < *num_retvals; i++) {
// TODO(srbs): Manage refcount of ForwardOperation's inputs/outputs.
forward_op_.outputs.push_back(retvals[i]);
}
// TODO(b/166669239): This is needed to support AttrBuilder::Get for string
// attributes. Number type attrs and DataType attrs work fine without this.
// Consider getting rid of this and making the behavior between number types
// and string consistent.
forward_op_.attrs.BuildNodeDef();
// TODO(b/170307493): Populate skip_input_indices here.
std::unique_ptr<GradientFunction> backward_fn;
TF_RETURN_IF_ERROR(registry_.Lookup(forward_op_, &backward_fn));
tape_->RecordOperation(forward_op_.inputs, forward_op_.outputs,
backward_fn.release(), parent_op_->Name());
return absl::OkStatus();
}
} // namespace gradients
} // namespace tensorflow
tensorflow/c/experimental/gradients/tape/tape_operation.h
+94
View File
@@ -0,0 +1,94 @@
/* 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.
==============================================================================*/
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_
#include <cstddef>
#include <cstdint>
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_operation.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/gradients.h"
#include "tensorflow/c/tensor_interface.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace gradients {
class TapeOperation : public AbstractOperation {
public:
explicit TapeOperation(AbstractOperation*, Tape*, const GradientRegistry&);
void Release() override;
absl::Status Reset(const char* op, const char* raw_device_name) override;
const std::string& Name() const override;
const std::string& DeviceName() const override;
absl::Status SetDeviceName(const char* name) override;
absl::Status AddInput(AbstractTensorHandle* input) override;
absl::Status AddInputList(
absl::Span<AbstractTensorHandle* const> inputs) override;
absl::Status Execute(absl::Span<AbstractTensorHandle*> retvals,
int* num_retvals) override;
absl::Status SetAttrString(const char* attr_name, const char* data,
size_t length) override;
absl::Status SetAttrInt(const char* attr_name, int64_t value) override;
absl::Status SetAttrFloat(const char* attr_name, float value) override;
absl::Status SetAttrBool(const char* attr_name, bool value) override;
absl::Status SetAttrType(const char* attr_name, DataType value) override;
absl::Status SetAttrShape(const char* attr_name, const int64_t* dims,
const int num_dims) override;
absl::Status SetAttrFunction(const char* attr_name,
const AbstractOperation* value) override;
absl::Status SetAttrFunctionName(const char* attr_name, const char* value,
size_t length) override;
absl::Status SetAttrTensor(const char* attr_name,
AbstractTensorInterface* tensor) override;
absl::Status SetAttrStringList(const char* attr_name,
const void* const* values,
const size_t* lengths,
int num_values) override;
absl::Status SetAttrFloatList(const char* attr_name, const float* values,
int num_values) override;
absl::Status SetAttrIntList(const char* attr_name, const int64_t* values,
int num_values) override;
absl::Status SetAttrTypeList(const char* attr_name, const DataType* values,
int num_values) override;
absl::Status SetAttrBoolList(const char* attr_name,
const unsigned char* values,
int num_values) override;
absl::Status SetAttrShapeList(const char* attr_name, const int64_t** dims,
const int* num_dims, int num_values) override;
absl::Status SetAttrFunctionList(
const char* attr_name,
absl::Span<const AbstractOperation*> values) override;
AbstractOperation* GetBackingOperation();
// For LLVM style RTTI.
static bool classof(const AbstractOperation* ptr) {
return ptr->getKind() == kTape;
}
~TapeOperation() override;
private:
AbstractOperation* parent_op_;
ForwardOperation forward_op_;
Tape* tape_;
const GradientRegistry& registry_;
};
} // namespace gradients
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRADIENTS_TAPE_TAPE_OPERATION_H_
tensorflow/c/experimental/grappler/BUILD
+77
View File
@@ -0,0 +1,77 @@
# Description:
# Graph C API.
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "grappler_hdrs",
hdrs = ["grappler.h"],
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/c:c_api",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_status_headers",
],
)
cc_library(
name = "grappler",
srcs = ["grappler.cc"],
hdrs = [
"grappler.h",
"grappler_internal.h",
],
visibility = ["//tensorflow:internal"],
deps = [
"//tensorflow/c:c_api_internal",
"//tensorflow/c:c_api_macros",
"//tensorflow/c:tf_buffer",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status",
"//tensorflow/c:tf_status_helper",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/costs:graph_properties",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
"//tensorflow/core/platform:logging",
"//tensorflow/core/platform:status",
"@com_google_absl//absl/status",
"@xla//xla/tsl/platform:env",
"@xla//xla/tsl/platform:errors",
],
)
tf_cc_test(
name = "grappler_test",
srcs = ["grappler_test.cc"],
deps = [
":grappler",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status_headers",
"//tensorflow/core:framework",
"//tensorflow/core:portable_gif_internal",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler/clusters:single_machine",
"//tensorflow/core/grappler/costs:op_performance_data_cc",
"//tensorflow/core/grappler/inputs:trivial_test_graph_input_yielder",
"//tensorflow/core/grappler/optimizers:custom_graph_optimizer_registry",
"//tensorflow/core/platform:status",
"//tensorflow/core/protobuf:for_core_protos_cc",
"@com_google_absl//absl/log:check",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tensorflow/c/experimental/grappler/grappler.cc
+396
View File
@@ -0,0 +1,396 @@
/* Copyright 2021 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.
==============================================================================*/
// This file extends/implements core graph optimizer base classes in terms of
// the C API defined in grappler.h. A class "CSomething" represents a
// "Something" that can be manipulated via calls in the C interface and a C
// struct called "TP_Something".
#include "tensorflow/c/experimental/grappler/grappler.h"
#include <algorithm>
#include <cstddef>
#include <cstring>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "absl/status/status.h"
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/experimental/grappler/grappler_internal.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_helper.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace {
#define VALIDATE_STRUCT_SIZE(STRUCT_NAME, STRUCT_OBJ, SIZE_VALUE_NAME) \
do { \
if (STRUCT_OBJ.struct_size == 0) { \
return absl::Status(absl::StatusCode::kFailedPrecondition, \
"struct_size field in " #STRUCT_NAME \
" must be set to " #SIZE_VALUE_NAME "."); \
} \
} while (0)
#define VALIDATE_MEMBER(STRUCT_NAME, STRUCT_OBJ, NAME) \
do { \
if (STRUCT_OBJ.NAME == 0) { \
return absl::Status(absl::StatusCode::kFailedPrecondition, \
"'" #NAME "' field in " #STRUCT_NAME \
" must be set."); \
} \
} while (0)
absl::Status ValidateTPOptimizerRegistrationParams(
const TP_OptimizerRegistrationParams& params) {
VALIDATE_STRUCT_SIZE(TP_OptimizerRegistrationParams, params,
TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE);
VALIDATE_MEMBER(TP_OptimizerRegistrationParams, params, device_type);
return absl::OkStatus();
}
absl::Status ValidateTPOptimizer(const TP_Optimizer& optimizer) {
VALIDATE_STRUCT_SIZE(TP_Optimizer, optimizer, TP_OPTIMIZER_STRUCT_SIZE);
VALIDATE_MEMBER(TP_Optimizer, optimizer, optimize_func);
return absl::OkStatus();
}
absl::Status ValidateTPOptimizerConfigs(const TP_OptimizerConfigs& configs) {
VALIDATE_STRUCT_SIZE(TP_OptimizerConfigs, configs,
TP_OPTIMIZER_CONFIGS_STRUCT_SIZE);
return absl::OkStatus();
}
#undef VALIDATE_MEMBER
#undef VALIDATE_STRUCT_SIZE
} // namespace
namespace tensorflow {
namespace grappler {
absl::Status CGraphOptimizer::Optimize(Cluster* cluster,
const GrapplerItem& item,
GraphDef* optimized_graph_def) {
OwnedTFStatus c_status(TF_NewStatus());
OwnedTFBuffer graph_buf(TF_NewBuffer());
OwnedTFBuffer optimized_graph_buf(TF_NewBuffer());
TF_RETURN_IF_ERROR(MessageToBuffer(item.graph, graph_buf.get()));
optimizer_.optimize_func(c_optimizer_, graph_buf.get(),
reinterpret_cast<const TF_GrapplerItem*>(&item),
optimized_graph_buf.get(), c_status.get());
TF_RETURN_IF_ERROR(tsl::StatusFromTF_Status(c_status.get()));
TF_RETURN_IF_ERROR(
BufferToMessage(optimized_graph_buf.get(), optimized_graph_def));
return absl::OkStatus();
}
#define CONFIG_TOGGLE(optimizer) \
if (tp_configs.optimizer == TF_TriState_Off) \
configs.toggle_config[#optimizer] = RewriterConfig::OFF; \
else \
configs.toggle_config[#optimizer] = RewriterConfig::ON;
void CGraphOptimizerRegister(
const PluginGraphOptimizerRegistry::Creator& creator,
const TP_OptimizerConfigs tp_configs, const char* device_type) {
ConfigList configs;
// disable_model_pruning is turned off by default.
if (tp_configs.disable_model_pruning == TF_TriState_On)
configs.disable_model_pruning = true;
else
configs.disable_model_pruning = false;
// The other configs are turned on by default.
CONFIG_TOGGLE(implementation_selector);
CONFIG_TOGGLE(function_optimization);
CONFIG_TOGGLE(common_subgraph_elimination);
CONFIG_TOGGLE(arithmetic_optimization);
CONFIG_TOGGLE(debug_stripper);
CONFIG_TOGGLE(constant_folding);
CONFIG_TOGGLE(shape_optimization);
CONFIG_TOGGLE(auto_mixed_precision);
CONFIG_TOGGLE(auto_mixed_precision_onednn_bfloat16);
CONFIG_TOGGLE(auto_mixed_precision_mkl);
CONFIG_TOGGLE(pin_to_host_optimization);
CONFIG_TOGGLE(layout_optimizer);
CONFIG_TOGGLE(remapping);
CONFIG_TOGGLE(loop_optimization);
CONFIG_TOGGLE(dependency_optimization);
CONFIG_TOGGLE(auto_parallel);
CONFIG_TOGGLE(memory_optimization);
CONFIG_TOGGLE(scoped_allocator_optimization);
PluginGraphOptimizerRegistry::RegisterPluginOptimizerOrDie(
creator, device_type, configs);
}
#undef CONFIG_TOGGLE
absl::Status InitGraphPlugin(void* dso_handle) {
tsl::Env* env = tsl::Env::Default();
// Step 1: Load symbol for `TF_InitPlugin`
void* dso_symbol;
TF_RETURN_IF_ERROR(
env->GetSymbolFromLibrary(dso_handle, "TF_InitGraph", &dso_symbol));
// Step 2: Call `TF_InitPlugin`
auto init_fn = reinterpret_cast<TFInitGraphPluginFn>(dso_symbol);
return InitGraphPlugin(init_fn);
}
absl::Status InitGraphPlugin(TFInitGraphPluginFn init_fn) {
TP_OptimizerRegistrationParams params{
TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE};
TP_Optimizer optimizer{TP_OPTIMIZER_STRUCT_SIZE};
TP_OptimizerConfigs optimizer_configs{TP_OPTIMIZER_CONFIGS_STRUCT_SIZE};
params.major_version = GO_MAJOR;
params.minor_version = GO_MINOR;
params.patch_version = GO_PATCH;
params.optimizer = &optimizer;
params.optimizer_configs = &optimizer_configs;
OwnedTFStatus c_status(TF_NewStatus());
init_fn(&params, c_status.get());
TF_RETURN_IF_ERROR(tsl::StatusFromTF_Status(c_status.get()));
TF_RETURN_IF_ERROR(ValidateTPOptimizerRegistrationParams(params));
TF_RETURN_IF_ERROR(ValidateTPOptimizer(optimizer));
TF_RETURN_IF_ERROR(ValidateTPOptimizerConfigs(optimizer_configs));
CGraphOptimizerRegister(
[=]() { return new CGraphOptimizer(optimizer, params.device_type); },
optimizer_configs, params.device_type);
return absl::OkStatus();
}
} // namespace grappler
} // namespace tensorflow
void TF_GetNodesToPreserveListSize(const TF_GrapplerItem* item, int* num_values,
size_t* storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::unordered_set<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)
->NodesToPreserve();
*num_values = nodes.size();
*storage_size = 0;
for (const std::string& str : nodes) {
*storage_size += str.size();
}
}
void TF_GetNodesToPreserveList(const TF_GrapplerItem* item, char** values,
size_t* lengths, int num_values, void* storage,
size_t storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::unordered_set<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)
->NodesToPreserve();
char* p = static_cast<char*>(storage);
int index = 0;
for (const std::string& s : nodes) {
if (index >= num_values) break;
values[index] = p;
lengths[index] = s.size();
if ((p + s.size()) > (static_cast<char*>(storage) + storage_size)) {
tsl::Set_TF_Status_from_Status(
status,
absl::InvalidArgumentError(
"Not enough storage to hold the requested list of nodes"));
return;
}
memcpy(values[index], s.data(), s.size());
p += s.size();
index++;
}
}
void TF_GetFetchNodesListSize(const TF_GrapplerItem* item, int* num_values,
size_t* storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)->fetch;
*num_values = nodes.size();
*storage_size = 0;
for (const std::string& str : nodes) {
*storage_size += str.size();
}
}
void TF_GetFetchNodesList(const TF_GrapplerItem* item, char** values,
size_t* lengths, int num_values, void* storage,
size_t storage_size, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<std::string>& nodes =
reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)->fetch;
const int len = std::min(num_values, static_cast<int>(nodes.size()));
char* p = static_cast<char*>(storage);
for (int index = 0; index < len; ++index) {
const std::string& s = nodes[index];
values[index] = p;
lengths[index] = s.size();
if ((p + s.size()) > (static_cast<char*>(storage) + storage_size)) {
tsl::Set_TF_Status_from_Status(
status,
absl::InvalidArgumentError(
"Not enough storage to hold the requested list of nodes"));
return;
}
memcpy(values[index], s.data(), s.size());
p += s.size();
}
}
TF_GraphProperties* TF_NewGraphProperties(const TF_GrapplerItem* item) {
return reinterpret_cast<TF_GraphProperties*>(
new tensorflow::grappler::GraphProperties(
*reinterpret_cast<const tensorflow::grappler::GrapplerItem*>(item)));
}
void TF_DeleteGraphProperties(TF_GraphProperties* graph_properties) {
if (graph_properties == nullptr) return;
delete reinterpret_cast<tensorflow::grappler::GraphProperties*>(
graph_properties);
}
void TF_InferStatically(TF_GraphProperties* graph_properties,
TF_Bool assume_valid_feeds,
TF_Bool aggressive_shape_inference,
TF_Bool include_input_tensor_values,
TF_Bool include_output_tensor_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
absl::Status s =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->InferStatically(assume_valid_feeds, aggressive_shape_inference,
include_input_tensor_values,
include_output_tensor_values);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
}
}
void TF_GetInputPropertiesListSize(TF_GraphProperties* graph_properties,
const char* name, int* num_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
*num_values =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetInputProperties(name)
.size();
}
void TF_GetOutputPropertiesListSize(TF_GraphProperties* graph_properties,
const char* name, int* num_values,
TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
*num_values =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetOutputProperties(name)
.size();
}
void TF_GetInputPropertiesList(TF_GraphProperties* graph_properties,
const char* name, TF_Buffer** properties,
int num_values, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<tensorflow::OpInfo::TensorProperties>& tensor_properties =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetInputProperties(name);
const int len =
std::min(num_values, static_cast<int>(tensor_properties.size()));
for (int i = 0; i < len; ++i) {
absl::Status s =
tensorflow::MessageToBuffer(tensor_properties[i], properties[i]);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
}
void TF_GetOutputPropertiesList(TF_GraphProperties* graph_properties,
const char* name, TF_Buffer** properties,
int num_values, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const std::vector<tensorflow::OpInfo::TensorProperties>& tensor_properties =
reinterpret_cast<tensorflow::grappler::GraphProperties*>(graph_properties)
->GetOutputProperties(name);
const int len =
std::min(num_values, static_cast<int>(tensor_properties.size()));
for (int i = 0; i < len; ++i) {
absl::Status s =
tensorflow::MessageToBuffer(tensor_properties[i], properties[i]);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
}
TF_FunctionLibraryDefinition* TF_NewFunctionLibraryDefinition(
const TF_Buffer* graph_buf, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
tensorflow::GraphDef graph_def;
absl::Status s = tensorflow::BufferToMessage(graph_buf, &graph_def);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return nullptr;
}
return reinterpret_cast<TF_FunctionLibraryDefinition*>(
new tensorflow::FunctionLibraryDefinition(
tensorflow::OpRegistry::Global(), graph_def.library()));
}
void TF_DeleteFunctionLibraryDefinition(TF_FunctionLibraryDefinition* fn_lib) {
if (fn_lib == nullptr) return;
delete reinterpret_cast<tensorflow::FunctionLibraryDefinition*>(fn_lib);
}
void TF_LookUpOpDef(TF_FunctionLibraryDefinition* fn_lib, const char* name,
TF_Buffer* buf, TF_Status* status) {
TF_SetStatus(status, TF_OK, "");
const tensorflow::OpDef* op_def_ptr = nullptr;
absl::Status s =
reinterpret_cast<tensorflow::FunctionLibraryDefinition*>(fn_lib)
->LookUpOpDef(name, &op_def_ptr);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
s = tensorflow::MessageToBuffer(*op_def_ptr, buf);
if (!s.ok()) {
tsl::Set_TF_Status_from_Status(status, s);
return;
}
}
tensorflow/c/experimental/grappler/grappler.h
+294
View File
@@ -0,0 +1,294 @@
/* Copyright 2021 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_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
#include <stddef.h>
#include <stdint.h>
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
// --------------------------------------------------------------------------
// C API for Graph. The API is under active development and eventually
// should allow registering a plugin graph optimizer with TensorFlow.
//
// Conventions:
// * Struct prefix indicates whether struct fields should be filled by the
// plugin or core implementation:
// * Struct that should be filled by the plugin: `TP_OptimizerConfigs`,
// `TP_Optimizer`, `TP_OptimizerRegistrationParams`
// * Struct that should be filled by the proper: `TF_GrapplerItem`,
// `TF_GraphProperties`, `TF_FunctionLibraryDefinition`
// * We use `struct_size` for version checking. It should be set both by
// core and the plugin.
// * For example, `TF_InitGraph` function receives
// `TP_OptimizerRegistrationParams*` as input with `struct_size`
// populated by core. The plugin is responsible for setting
// `struct_size` as well, along with all other fields.
// * Refer to "TensorFlow Versioning Strategy" section at
// https://github.com/tensorflow/community/pull/257/files.
// * Note that the API is still under active development and doesn't have
// versioning guarantees yet.
// * `void* ext` is a free-form field that can be populated by
// a plugin in `TP_*` structs or potential future extension points .
//
// Example usage:
//
// /* Sample TensorFlow code below, exact implementation might differ. */
// // Version checking uses `struct_size`. It should be set both by core
// // and the plugin.
// TP_OptimizerRegistrationParams params{
// TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE};
// TP_Optimizer optimizer{TP_OPTIMIZER_STRUCT_SIZE};
// TP_OptimizerConfigs configs{TP_OPTIMIZER_CONFIGS_STRUCT_SIZE};
// params.optimizer = &optimizer;
// params.configs = &configs;
//
// /* Plugin code below */
// void TF_InitGraph(TP_OptimizerRegistrationParams* params,
// TF_Status* status) {
// params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE;
// params->device_type = "MY_DEVICE";
//
// // Disable certain optimizer.
// params->optimizer_configs->struct_size =
// TP_OPTIMIZER_CONFIGS_STRUCT_SIZE; params->optimizer_configs->remapping =
// TF_TriState_Off;
//
// // Set functions to create a new optimizer.
// params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE;
// params->optimizer->create_func = (My_optimizer::create_func);
// }
#define GO_MAJOR 0
#define GO_MINOR 0
#define GO_PATCH 1
#ifdef __cplusplus
extern "C" {
#endif
// TF_TriState is the C API typedef for tri-state.
typedef enum TF_TriState {
TF_TriState_Default = 0,
TF_TriState_Off,
TF_TriState_On,
} TF_TriState;
// TF_GrapplerItem represents a combination of a graph, one of more fetch nodes,
// and potentially a set of nodes to feed.
typedef struct TF_GrapplerItem TF_GrapplerItem;
// Flags indicating whether existing optimizers should be turned off.
// It's optional for plugin to set functions to return true/false. If not
// set, proper uses configuration set by user.
typedef struct TP_OptimizerConfigs {
size_t struct_size;
void* ext; // reserved for future use
TF_TriState disable_model_pruning;
TF_TriState implementation_selector;
TF_TriState function_optimization;
TF_TriState common_subgraph_elimination;
TF_TriState arithmetic_optimization;
TF_TriState debug_stripper;
TF_TriState constant_folding;
TF_TriState shape_optimization;
TF_TriState auto_mixed_precision;
TF_TriState auto_mixed_precision_onednn_bfloat16;
TF_TriState auto_mixed_precision_mkl;
TF_TriState pin_to_host_optimization;
TF_TriState layout_optimizer;
TF_TriState remapping;
TF_TriState loop_optimization;
TF_TriState dependency_optimization;
TF_TriState auto_parallel;
TF_TriState memory_optimization;
TF_TriState scoped_allocator_optimization;
} TP_OptimizerConfigs;
#define TP_OPTIMIZER_CONFIGS_STRUCT_SIZE \
TF_OFFSET_OF_END(TP_OptimizerConfigs, scoped_allocator_optimization)
// Struct for Optimizer. Plugin authors must provide an optimize function.
// Creation and deletion functions are optional.
typedef struct TP_Optimizer {
size_t struct_size;
void* ext; // reserved for future use
// [Optional]
// Create function for optimizer.
void* (*create_func)();
// Optimizer function for optimizer. The first param is an optimizer created
// by create_func. The second param is input graph. The third param is
// GrapplerItem. The fourth param is output graph.
void (*optimize_func)(void*, const TF_Buffer*, const TF_GrapplerItem*,
TF_Buffer*, TF_Status*);
// [Optional]
// Destroy function for optimizer. If Create function is provided, destroy
// function is must.
void (*destroy_func)(void*);
} TP_Optimizer;
#define TP_OPTIMIZER_STRUCT_SIZE TF_OFFSET_OF_END(TP_Optimizer, destroy_func)
typedef struct TP_OptimizerRegistrationParams {
size_t struct_size;
void* ext; // reserved for future use
// Graph C API version.
int32_t major_version;
int32_t minor_version;
int32_t patch_version;
// Backend device type supported by the optimizer.
const char* device_type;
TP_OptimizerConfigs* optimizer_configs; // output, set by plugin
TP_Optimizer* optimizer; // output, set by plugin
} TP_OptimizerRegistrationParams;
#define TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE \
TF_OFFSET_OF_END(TP_OptimizerRegistrationParams, optimizer)
// TF_InitGraph is used to do graph optimizer registration.
// Plugin should implement TF_InitGraph to register graph optimizers.
void TF_InitGraph(TP_OptimizerRegistrationParams* params, TF_Status* status);
// Get a set of node names that must be preserved. They can not be transformed
// or removed during the graph transformation. This includes feed and fetch
// nodes, keep_ops, init_ops. Fills in `num_values` and `storage_size`, they
// will be used in `TF_GetNodesToPreserveList`.
TF_CAPI_EXPORT extern void TF_GetNodesToPreserveListSize(
const TF_GrapplerItem* item, int* num_values, size_t* storage_size,
TF_Status* status);
// Get a set of node names that must be preserved. They can not be transformed
// or removed during the graph transformation. This includes feed and fetch
// nodes, keep_ops, init_ops. Fills in `values` and `lengths`, each of which
// must point to an array of length at least `num_values`.
//
// The elements of values will point to addresses in `storage` which must be at
// least `storage_size` bytes in length. `num_values` and `storage` can be
// obtained from TF_GetNodesToPreserveSize
//
// Fails if storage_size is too small to hold the requested number of strings.
TF_CAPI_EXPORT extern void TF_GetNodesToPreserveList(
const TF_GrapplerItem* item, char** values, size_t* lengths, int num_values,
void* storage, size_t storage_size, TF_Status* status);
// Get a set of node names for fetch nodes. Fills in `values` and `lengths`,
// they will be used in `TF_GetFetchNodesList`
TF_CAPI_EXPORT extern void TF_GetFetchNodesListSize(const TF_GrapplerItem* item,
int* num_values,
size_t* storage_size,
TF_Status* status);
// Get a set of node names for fetch nodes. Fills in `values` and `lengths`,
// each of which must point to an array of length at least `num_values`.
//
// The elements of values will point to addresses in `storage` which must be at
// least `storage_size` bytes in length. `num_values` and `storage` can be
// obtained from TF_GetFetchNodesSize
//
// Fails if storage_size is too small to hold the requested number of strings.
TF_CAPI_EXPORT extern void TF_GetFetchNodesList(const TF_GrapplerItem* item,
char** values, size_t* lengths,
int num_values, void* storage,
size_t storage_size,
TF_Status* status);
// Infer OpInfo::TensorProperties for graph nodes inputs/outputs.
//
// Typical use case, is to infer tensor properties from a graph, before doing
// optimization pass. Nodes modified during optimization pass have to be
// invalidated, to prevent further incorrect optimizations based on wrong shape
// and data type properties.
typedef struct TF_GraphProperties TF_GraphProperties;
// Create GraphProperties. The item must outlive the properties.
TF_CAPI_EXPORT extern TF_GraphProperties* TF_NewGraphProperties(
const TF_GrapplerItem* item);
// Delete GraphProperties.
TF_CAPI_EXPORT extern void TF_DeleteGraphProperties(
TF_GraphProperties* graph_properties);
// Infer tensor shapes through abstract interpretation.
// If assume_valid_feeds is true, it can help infer shapes in the fanout of fed
// nodes. This may cause incorrectness in graph analyses, but is useful for
// simulation or scheduling.
// If aggressive_shape_inference is true, nodes are executed on the host to
// identify output values when possible and does other aggressive strategies.
// This may cause incorrectness in graph analyses, but is useful for simulation
// or scheduling.
// If include_input_tensor_values is true, the values of constant
// tensors will included in the input properties.
// If include_output_tensor_values is true, the values of constant tensors will
// be included in the output properties.
TF_CAPI_EXPORT extern void TF_InferStatically(
TF_GraphProperties* graph_properties, TF_Bool assume_valid_feeds,
TF_Bool aggressive_shape_inference, TF_Bool include_input_tensor_values,
TF_Bool include_output_tensor_values, TF_Status* s);
// Get the size of input OpInfo::TensorProperties given node name.
TF_CAPI_EXPORT extern void TF_GetInputPropertiesListSize(
TF_GraphProperties* graph_properties, const char* name, int* num_values,
TF_Status* status);
// Get the size of output OpInfo::TensorProperties given node name.
TF_CAPI_EXPORT extern void TF_GetOutputPropertiesListSize(
TF_GraphProperties* graph_properties, const char* name, int* num_values,
TF_Status* status);
// Get a list of input OpInfo::TensorProperties given node name.
// Return the serialized list `properties`.
TF_CAPI_EXPORT extern void TF_GetInputPropertiesList(
TF_GraphProperties* graph_properties, const char* name,
TF_Buffer** properties, int num_values, TF_Status* status);
// Get a list of output OpInfo::TensorProperties given node name.
// Return the serialized list `properties`.
TF_CAPI_EXPORT extern void TF_GetOutputPropertiesList(
TF_GraphProperties* graph_properties, const char* name,
TF_Buffer** properties, int num_values, TF_Status* status);
// Helper to maintain a map between function names in a given
// FunctionDefLibrary and function definitions.
// Typical use case, is to look up an OpDef by type name.
typedef struct TF_FunctionLibraryDefinition TF_FunctionLibraryDefinition;
// Create NewFunctionLibraryDefinition.
TF_CAPI_EXPORT extern TF_FunctionLibraryDefinition*
TF_NewFunctionLibraryDefinition(const TF_Buffer* graph_buf, TF_Status* status);
// Delete NewFunctionLibraryDefinition.
TF_CAPI_EXPORT extern void TF_DeleteFunctionLibraryDefinition(
TF_FunctionLibraryDefinition* fn_lib);
// Shorthand for calling LookUp to get the OpDef from FunctionLibraryDefinition
// given op name. The returned OpDef is represented by TF_Buffer.
TF_CAPI_EXPORT extern void TF_LookUpOpDef(TF_FunctionLibraryDefinition* fn_lib,
const char* name, TF_Buffer* buf,
TF_Status* s);
#ifdef __cplusplus
} // extern "C"
#endif
#endif // TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_H_
tensorflow/c/experimental/grappler/grappler_internal.h
+104
View File
@@ -0,0 +1,104 @@
/* Copyright 2021 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.
==============================================================================*/
// Classes and utilities that work with Graph C API for internal use.
// This includes functions used for optimizer registration and interfaces needed
// for testing.
#ifndef TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
#define TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
#include <functional>
#include <memory>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#include "tensorflow/c/c_api.h"
#include "tensorflow/c/experimental/grappler/grappler.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
// Plugin initialization function that a device plugin
// must define.
typedef void (*TFInitGraphPluginFn)(TP_OptimizerRegistrationParams* const,
TF_Status* const);
// Registers Graph optimizers.
Status InitGraphPlugin(void* dso_handle);
// Allow registering a graph optimizer using a function (used for
// testing).
Status InitGraphPlugin(TFInitGraphPluginFn init_fn);
struct GrapplerItem;
class Cluster;
struct TFStatusDeleter {
void operator()(TF_Status* s) const { TF_DeleteStatus(s); }
};
using OwnedTFStatus = std::unique_ptr<TF_Status, TFStatusDeleter>;
struct TFBufferDeleter {
void operator()(TF_Buffer* buf) const { TF_DeleteBuffer(buf); }
};
using OwnedTFBuffer = std::unique_ptr<TF_Buffer, TFBufferDeleter>;
class CGraphOptimizer : public CustomGraphOptimizer {
public:
explicit CGraphOptimizer(TP_Optimizer optimizer, const char* device_type)
: optimizer_(optimizer), device_type_(device_type) {
if (optimizer.create_func != nullptr) {
c_optimizer_ = (*optimizer_.create_func)();
} else {
c_optimizer_ = nullptr;
}
}
std::string name() const override { return "PluggableGraphOptimizer"; }
bool UsesFunctionLibrary() const override { return false; }
Status Init(
const tensorflow::RewriterConfig_CustomGraphOptimizer* config) override {
return OkStatus();
}
Status Optimize(Cluster* cluster, const GrapplerItem& item,
GraphDef* optimized_graph_def) override;
~CGraphOptimizer() override {
if (optimizer_.destroy_func != nullptr) {
(*optimizer_.destroy_func)(c_optimizer_);
}
}
private:
TP_Optimizer optimizer_;
std::string device_type_;
void* c_optimizer_;
};
// Registration function to register a CGraphOptimizer along with plugin configs
// and device type.
void CGraphOptimizerRegister(
const PluginGraphOptimizerRegistry::Creator& creator,
const TP_OptimizerConfigs tp_configs, const char* device_type);
} // namespace grappler
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_GRAPPLER_GRAPPLER_INTERNAL_H_
tensorflow/c/experimental/grappler/grappler_test.cc
+329
View File
@@ -0,0 +1,329 @@
/* Copyright 2021 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/c/experimental/grappler/grappler.h"
#include <cstddef>
#include <memory>
#include <set>
#include <string>
#include <unordered_set>
#include <vector>
#include "absl/log/check.h"
#include "tensorflow/c/experimental/grappler/grappler_internal.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_buffer_internal.h"
#include "tensorflow/c/tf_status.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "tensorflow/core/framework/function.h"
#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/grappler/clusters/single_machine.h"
#include "tensorflow/core/grappler/costs/op_performance_data.pb.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/inputs/trivial_test_graph_input_yielder.h"
#include "tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
namespace {
void optimize_func(void* optimizer, const TF_Buffer* graph_buf,
const TF_GrapplerItem* item, TF_Buffer* optimized_graph_buf,
TF_Status* tf_status) {}
void PopulateDefaultParam(TP_OptimizerRegistrationParams* params) {
params->struct_size = TP_OPTIMIZER_REGISTRATION_PARAMS_STRUCT_SIZE;
params->optimizer_configs->struct_size = TP_OPTIMIZER_CONFIGS_STRUCT_SIZE;
params->optimizer->struct_size = TP_OPTIMIZER_STRUCT_SIZE;
params->optimizer->create_func = nullptr;
params->optimizer->optimize_func = optimize_func;
params->optimizer->destroy_func = nullptr;
}
TEST(Grappler, SuccessfulRegistration) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Success";
params->optimizer_configs->remapping = TF_TriState_Off;
};
TF_ASSERT_OK(InitGraphPlugin(plugin_init));
ASSERT_EQ(PluginGraphOptimizerRegistry::CreateOptimizers(
std::set<std::string>{"Success"})
.size(),
1);
ConfigList config = PluginGraphOptimizerRegistry::GetPluginConfigs(
true, std::set<std::string>{"Success"});
ASSERT_EQ(config.toggle_config["remapping"], RewriterConfig::OFF);
}
TEST(Grappler, MultiplePluginRegistration) {
auto plugin_init_0 = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Device0";
};
auto plugin_init_1 = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "Device1";
};
TF_ASSERT_OK(InitGraphPlugin(plugin_init_0));
TF_ASSERT_OK(InitGraphPlugin(plugin_init_1));
ASSERT_EQ(PluginGraphOptimizerRegistry::CreateOptimizers(
std::set<std::string>{"Device0", "Device1"})
.size(),
2);
}
TEST(Grappler, DeviceTypeNotSet) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = nullptr;
};
absl::Status status = InitGraphPlugin(plugin_init);
ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
ASSERT_EQ(
status.message(),
"'device_type' field in TP_OptimizerRegistrationParams must be set.");
}
TEST(Grappler, OptimizeFuncNotSet) {
auto plugin_init = [](TP_OptimizerRegistrationParams* const params,
TF_Status* const status) -> void {
TF_SetStatus(status, TF_OK, "");
PopulateDefaultParam(params);
params->device_type = "FuncNotSet";
params->optimizer->optimize_func = nullptr;
};
absl::Status status = InitGraphPlugin(plugin_init);
ASSERT_EQ(status.code(), tensorflow::error::FAILED_PRECONDITION);
ASSERT_EQ(status.message(),
"'optimize_func' field in TP_Optimizer must be set.");
}
TEST(TF_GrapplerItem, NodesToPreserve) {
GrapplerItem item;
item.fetch = std::vector<std::string>{"Conv", "BiasAdd"};
std::unordered_set<std::string> nodes_preserved = item.NodesToPreserve();
TF_GrapplerItem* c_item = reinterpret_cast<TF_GrapplerItem*>(&item);
int list_total_size = 0;
for (const std::string& s : nodes_preserved) {
list_total_size += s.size();
}
size_t storage_size = 0;
int num_values = 0;
TF_Status* status = TF_NewStatus();
TF_GetNodesToPreserveListSize(c_item, &num_values, &storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(nodes_preserved.size(), num_values);
EXPECT_EQ(list_total_size, storage_size);
std::unique_ptr<char*[]> values(new char*[nodes_preserved.size()]);
std::unique_ptr<size_t[]> lens(new size_t[nodes_preserved.size()]);
std::unique_ptr<char[]> storage(new char[storage_size]);
TF_GetNodesToPreserveList(c_item, values.get(), lens.get(),
nodes_preserved.size(), storage.get(), storage_size,
status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (size_t i = 0; i < nodes_preserved.size(); ++i) {
EXPECT_EQ(
nodes_preserved.find(std::string(static_cast<const char*>(values[i]),
lens[i])) != nodes_preserved.end(),
true);
}
TF_DeleteStatus(status);
}
TEST(TF_GrapplerItem, FetchNodes) {
GrapplerItem item;
item.fetch = std::vector<std::string>{"Conv", "BiasAdd"};
TF_GrapplerItem* c_item = reinterpret_cast<TF_GrapplerItem*>(&item);
int list_total_size = 0;
for (const std::string& s : item.fetch) {
list_total_size += s.size();
}
size_t storage_size = 0;
int num_values = 0;
TF_Status* status = TF_NewStatus();
TF_GetFetchNodesListSize(c_item, &num_values, &storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(item.fetch.size(), num_values);
EXPECT_EQ(list_total_size, storage_size);
std::unique_ptr<char*[]> values(new char*[item.fetch.size()]);
std::unique_ptr<size_t[]> lens(new size_t[item.fetch.size()]);
std::unique_ptr<char[]> storage(new char[storage_size]);
TF_GetFetchNodesList(c_item, values.get(), lens.get(), item.fetch.size(),
storage.get(), storage_size, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (size_t i = 0; i < item.fetch.size(); ++i) {
EXPECT_EQ(item.fetch[i].size(), lens[i]) << i;
EXPECT_EQ(item.fetch[i],
std::string(static_cast<const char*>(values[i]), lens[i]))
<< i;
}
TF_DeleteStatus(status);
}
TEST(TF_GraphProperties, InputProperties) {
std::unique_ptr<SingleMachine> cluster(new SingleMachine(5 * 60, 3, 0));
TF_ASSERT_OK(cluster->Provision());
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_Status* status = TF_NewStatus();
TF_GraphProperties* graph_properties =
TF_NewGraphProperties(reinterpret_cast<TF_GrapplerItem*>(&item));
TF_InferStatically(graph_properties, true, false, false, false, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (const NodeDef& node : item.graph.node()) {
if (node.op() == "AddN") {
int num_values = 0;
TF_GetInputPropertiesListSize(graph_properties, node.name().c_str(),
&num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(num_values, 1);
std::vector<TF_Buffer*> in_props_buf(num_values, TF_NewBuffer());
TF_GetInputPropertiesList(graph_properties, node.name().c_str(),
in_props_buf.data(), num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
tensorflow::OpInfo::TensorProperties in_props;
absl::Status s = tensorflow::BufferToMessage(in_props_buf[0], &in_props);
TF_ASSERT_OK(s);
EXPECT_EQ(DT_FLOAT, in_props.dtype());
EXPECT_FALSE(in_props.shape().unknown_rank());
EXPECT_EQ(2, in_props.shape().dim_size());
EXPECT_EQ(10, in_props.shape().dim(0).size());
EXPECT_EQ(1, in_props.shape().dim(1).size());
for (int i = 0; i < in_props_buf.size(); i++)
TF_DeleteBuffer(in_props_buf[i]);
}
}
TF_DeleteGraphProperties(graph_properties);
TF_DeleteStatus(status);
TF_ASSERT_OK(cluster->Shutdown());
}
TEST(TF_GraphProperties, OutputProperties) {
std::unique_ptr<SingleMachine> cluster(new SingleMachine(5 * 60, 3, 0));
TF_ASSERT_OK(cluster->Provision());
TrivialTestGraphInputYielder fake_input(4, 1, 10, false,
cluster->GetDeviceNames());
GrapplerItem item;
CHECK(fake_input.NextItem(&item));
TF_Status* status = TF_NewStatus();
TF_GraphProperties* graph_properties =
TF_NewGraphProperties(reinterpret_cast<TF_GrapplerItem*>(&item));
TF_InferStatically(graph_properties, true, false, false, false, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
for (const NodeDef& node : item.graph.node()) {
if (node.op() == "AddN") {
int num_values = 0;
TF_GetOutputPropertiesListSize(graph_properties, node.name().c_str(),
&num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
EXPECT_EQ(num_values, 1);
std::vector<TF_Buffer*> out_props_buf(num_values, TF_NewBuffer());
TF_GetOutputPropertiesList(graph_properties, node.name().c_str(),
out_props_buf.data(), num_values, status);
EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
tensorflow::OpInfo::TensorProperties out_props;
absl::Status s =
tensorflow::BufferToMessage(out_props_buf[0], &out_props);
TF_ASSERT_OK(s);
EXPECT_EQ(DT_FLOAT, out_props.dtype());
EXPECT_FALSE(out_props.shape().unknown_rank());
EXPECT_EQ(2, out_props.shape().dim_size());
EXPECT_EQ(10, out_props.shape().dim(0).size());
EXPECT_EQ(1, out_props.shape().dim(1).size());
for (int i = 0; i < out_props_buf.size(); i++)
TF_DeleteBuffer(out_props_buf[i]);
}
}
TF_DeleteStatus(status);
TF_DeleteGraphProperties(graph_properties);
TF_ASSERT_OK(cluster->Shutdown());
}
TEST(TF_FunctionLibraryDefinition, LookUpOpDef) {
TF_Buffer* g_buf = TF_NewBuffer();
TF_Buffer* op_buf = TF_NewBuffer();
TF_Status* status = TF_NewStatus();
GraphDef g_def;
absl::Status s = MessageToBuffer(g_def, g_buf);
TF_ASSERT_OK(s);
TF_FunctionLibraryDefinition* func =
TF_NewFunctionLibraryDefinition(g_buf, status);
TF_LookUpOpDef(func, "Add", op_buf, status);
std::string actual_string(reinterpret_cast<const char*>(op_buf->data),
op_buf->length);
ASSERT_EQ(TF_OK, TF_GetCode(status));
const OpDef* expected_op_def;
TF_ASSERT_OK(OpRegistry::Global()->LookUpOpDef("Add", &expected_op_def));
std::string expected_serialized;
expected_op_def->SerializeToString(&expected_serialized);
EXPECT_EQ(expected_serialized, actual_string);
TF_DeleteBuffer(g_buf);
TF_DeleteBuffer(op_buf);
TF_DeleteStatus(status);
TF_DeleteFunctionLibraryDefinition(func);
}
} // namespace
} // namespace grappler
} // namespace tensorflow
tensorflow/c/experimental/next_pluggable_device/BUILD
+111
View File
@@ -0,0 +1,111 @@
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "c_api",
srcs = ["c_api.cc"],
hdrs = ["c_api.h"],
visibility = ["//visibility:public"],
deps = [
":tensor_pjrt_buffer_util",
"//tensorflow/c:c_api_macros_hdrs",
"//tensorflow/c:kernels_experimental_hdrs",
"//tensorflow/c:kernels_hdrs",
"//tensorflow/c:tf_buffer",
"//tensorflow/c:tf_status_internal",
"//tensorflow/c:tf_tensor_internal",
"//tensorflow/compiler/jit:variable_info",
"//tensorflow/compiler/jit:variable_info_util",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/common_runtime/next_pluggable_device:plugin_resource",
"//tensorflow/core/platform:refcount",
"//tensorflow/core/platform:status",
"//tensorflow/core/tfrt/common:pjrt_util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/time",
"@com_google_absl//absl/types:span",
"@xla//xla/pjrt:pjrt_client",
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
"@xla//xla/pjrt/c:pjrt_c_api_helpers",
"@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
"@xla//xla/tsl/distributed_runtime/coordination:coordination_service_agent",
],
)
# Plugin should include this target to avoid linking the C API implementation.
cc_library(
name = "c_api_hdrs",
hdrs = ["c_api.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/c:c_api_macros_hdrs",
"//tensorflow/c:kernels_hdrs",
"//tensorflow/c:tf_buffer_internal",
"//tensorflow/c:tf_status_headers",
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
],
)
cc_library(
name = "tensor_pjrt_buffer_util",
srcs = ["tensor_pjrt_buffer_util.cc"],
hdrs = ["tensor_pjrt_buffer_util.h"],
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/jit:pjrt_tensor_buffer_util",
"//tensorflow/core:framework",
"//tensorflow/core/tfrt/common:async_value_tensor",
"//tensorflow/core/tfrt/common:global_state",
"//tensorflow/core/tfrt/common:pjrt_state",
"//tensorflow/core/tfrt/common:pjrt_util",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@xla//xla/pjrt:pjrt_client",
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
"@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
"@xla//xla/tsl/platform:errors",
"@xla//xla/tsl/platform:statusor",
],
)
tf_cc_test(
name = "tensor_pjrt_buffer_util_test",
srcs = ["tensor_pjrt_buffer_util_test.cc"],
visibility = ["//visibility:public"],
deps = [
":tensor_pjrt_buffer_util",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/tfrt/common:async_value_tensor",
"//tensorflow/core/tfrt/common:pjrt_util",
"@com_google_absl//absl/base",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status:status_matchers",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
"@tsl//tsl/platform:casts",
"@xla//xla:shape_util",
"@xla//xla:xla_data_proto_cc",
"@xla//xla/pjrt:pjrt_api",
"@xla//xla/pjrt/c:pjrt_c_api_cpu",
"@xla//xla/pjrt/c:pjrt_c_api_hdrs",
"@xla//xla/pjrt/c:pjrt_c_api_wrapper_impl",
"@xla//xla/pjrt/c_api_client:pjrt_c_api_client",
"@xla//xla/pjrt/plugin/xla_cpu:cpu_client_options",
"@xla//xla/pjrt/plugin/xla_cpu:xla_cpu_pjrt_client",
"@xla//xla/tsl/lib/core:status_test_util",
"@xla//xla/tsl/platform:status_matchers",
"@xla//xla/tsl/platform:statusor",
"@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
],
)
tensorflow/c/experimental/next_pluggable_device/c_api.cc
+351
View File
@@ -0,0 +1,351 @@
/* Copyright 2022 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/c/experimental/next_pluggable_device/c_api.h"
#include <cstdint>
#include <cstdlib>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/time/time.h"
#include "absl/types/span.h"
#include "tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
#include "tensorflow/c/kernels.h"
#include "tensorflow/c/kernels_experimental.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
#include "tensorflow/c/tf_status_internal.h"
#include "tensorflow/c/tf_tensor.h"
#include "tensorflow/c/tf_tensor_internal.h"
#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/variable_info_util.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c/pjrt_c_api_helpers.h"
#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/tsl/distributed_runtime/coordination/coordination_service_agent.h"
#include "tensorflow/core/common_runtime/next_pluggable_device/plugin_resource.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/resource_handle.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/refcount.h"
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
TF_Device* TF_GetDevice(TF_OpKernelContext* ctx) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
return reinterpret_cast<TF_Device*>(cc_ctx->device());
}
// -------------------------- Resource ---------------------------------------
void TF_CreatePluginResource(TF_OpKernelContext* ctx,
const char* container_name,
const char* plugin_resource_name,
void* plugin_resource, void (*delete_func)(void*),
TF_Status* status) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
tensorflow::PluginResource* cc_resource_ptr = new tensorflow::PluginResource(
plugin_resource, plugin_resource_name, delete_func);
auto cc_status =
cc_ctx->resource_manager()->Create<tensorflow::PluginResource>(
container_name, plugin_resource_name, cc_resource_ptr);
status->status = cc_status;
}
void TF_LookupOrCreatePluginResource(
TF_OpKernelContext* ctx, const char* container_name,
const char* plugin_resource_name, void** result_plugin_resource,
void* (*create_func)(void*), void* create_func_args,
void (*delete_func)(void*), TF_Status* status) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
auto* resource_mgr = cc_ctx->resource_manager();
tensorflow::core::RefCountPtr<tensorflow::PluginResource>
tf_plugin_resource_ptr;
tensorflow::PluginResource* tf_plugin_resource = nullptr;
auto cc_status = resource_mgr->LookupOrCreate<tensorflow::PluginResource>(
container_name, plugin_resource_name, &tf_plugin_resource,
[plugin_resource_name, create_func, create_func_args,
delete_func](tensorflow::PluginResource** new_resource) {
void* opaque_plugin_resource = create_func(create_func_args);
*new_resource = new tensorflow::PluginResource(
opaque_plugin_resource, plugin_resource_name, delete_func);
return absl::OkStatus();
});
if (cc_status.ok()) {
tf_plugin_resource_ptr.reset(tf_plugin_resource);
*result_plugin_resource = tf_plugin_resource_ptr->GetOpaquePluginResource();
} else {
*result_plugin_resource = nullptr;
}
status->status = cc_status;
}
// ------------------------- VariableInfo ------------------------------------
struct TF_VariableInfo {
TF_VariableInfo() = delete;
// TF_VariableInfo is constructed here by TensorFlow, and will be passed to
// plugin as a opaque pointer. Plugin will need to call C APIs below to
// operate on TF_VariableInfo (such as allocate temp tensor for the `var` held
// by the underlying tensorflow::VariableInfo.
TF_VariableInfo(int index, const std::string& name, tensorflow::Var* var) {
var_info = tensorflow::VariableInfo{index, name, var};
}
tensorflow::VariableInfo var_info{0, "", nullptr};
};
TF_VariableInfo* TF_CreateVariableInfoFromContext(TF_OpKernelContext* ctx,
int index,
TF_Status* status) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
const tensorflow::Tensor& arg_tensor = cc_ctx->input(index);
absl::Status cc_status;
if (arg_tensor.dtype() != tensorflow::DT_RESOURCE) {
cc_status = absl::InvalidArgumentError(
absl::StrCat("Trying to obtain resource handle from Input[", index,
"], which is not type DT_RESOURCE."));
status->status = cc_status;
return nullptr;
}
const tensorflow::ResourceHandle& handle =
arg_tensor.flat<tensorflow::ResourceHandle>()(0);
tensorflow::Var* variable;
cc_status = tensorflow::LookupResource(cc_ctx, handle, &variable);
return new TF_VariableInfo(index, handle.name(), variable);
}
void TF_LockVariableInfos(TF_VariableInfo** vars, int num_vars,
TF_Status* status) {
std::vector<tensorflow::VariableInfo*> variable_ptrs;
variable_ptrs.reserve(num_vars);
for (int i = 0; i < num_vars; ++i) {
variable_ptrs.push_back(&(vars[i]->var_info));
}
absl::Status cc_status = LockVariables(absl::MakeSpan(variable_ptrs));
status->status = cc_status;
}
void TF_AllocateTempForVariableInfo(TF_OpKernelContext* ctx,
TF_VariableInfo* var_info,
TF_Status* status) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
absl::Status cc_status;
if (var_info == nullptr) {
cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
status->status = cc_status;
return;
}
if (var_info->var_info.var() == nullptr) {
cc_status = absl::InvalidArgumentError(
"VariableInfo does not track a resource variable.");
status->status = cc_status;
return;
}
cc_status = cc_ctx->allocate_temp(var_info->var_info.var()->tensor()->dtype(),
var_info->var_info.var()->tensor()->shape(),
var_info->var_info.var()->tensor());
status->status = cc_status;
}
TF_Tensor* TF_GetTensorFromVariableInfo(TF_VariableInfo* var_info,
TF_Status* status) {
absl::Status cc_status;
if (var_info == nullptr) {
cc_status = absl::InvalidArgumentError("TF_VariableInfo is NULL.");
status->status = cc_status;
return nullptr;
}
if (var_info->var_info.var() == nullptr) {
cc_status = absl::InvalidArgumentError(
"VariableInfo does not track a resource variable.");
status->status = cc_status;
return nullptr;
}
tensorflow::Tensor* tensor = var_info->var_info.var()->tensor();
TF_Tensor* result_tensor =
tensorflow::TF_TensorFromTensor(*tensor, &cc_status);
status->status = cc_status;
return result_tensor;
}
void TF_DeleteVariableInfo(TF_VariableInfo* var_info) {
if (var_info != nullptr) {
delete var_info;
}
}
// --------------------- Coordination service --------------------------------
TF_CoordinationServiceAgent* TF_GetCoordinationServiceAgent(
TF_OpKernelContext* ctx) {
auto* cc_ctx = reinterpret_cast<tensorflow::OpKernelContext*>(ctx);
return reinterpret_cast<TF_CoordinationServiceAgent*>(
cc_ctx->coordination_service_agent());
}
bool TF_CoordinationServiceIsInitialized(TF_CoordinationServiceAgent* agent) {
if (agent == nullptr) return false;
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
return cc_agent->IsInitialized();
}
void TF_CoordinationServiceInsertKeyValue(const char* key, int64_t key_size,
const char* value, int64_t value_size,
TF_CoordinationServiceAgent* agent,
TF_Status* status) {
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
absl::Status cc_status = cc_agent->InsertKeyValue(
absl::string_view(key, key_size), absl::string_view(value, value_size));
status->status = cc_status;
}
TF_Buffer* ProcessGetKeyValueResult(absl::StatusOr<std::string> value,
TF_Status* status) {
status->status = value.status();
if (!value.ok()) {
return nullptr;
}
// Caller is responsible to call `TF_DeleteBuffer` to release the buffer.
TF_Buffer* result = TF_NewBuffer();
const std::string& value_str = *value;
void* data = malloc(value_str.length());
value_str.copy(static_cast<char*>(data), value_str.length(), 0);
result->data = data;
result->length = value_str.length();
result->data_deallocator = [](void* data, size_t length) { free(data); };
return result;
}
TF_Buffer* TF_CoordinationServiceGetKeyValue(const char* key, int64_t key_size,
TF_CoordinationServiceAgent* agent,
TF_Status* status) {
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
auto value = cc_agent->GetKeyValue(absl::string_view(key, key_size));
return ProcessGetKeyValueResult(value, status);
}
TF_Buffer* TF_CoordinationServiceGetKeyValueWithTimeout(
const char* key, int64_t key_size, int64_t timeout_seconds,
TF_CoordinationServiceAgent* agent, TF_Status* status) {
if (timeout_seconds <= 0) {
status->status = absl::InvalidArgumentError(
"TF_CoordinationServiceGetKeyValueWithTimeout invoked with invalid "
"timeout_seconds <= 0.");
return nullptr;
}
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
auto value = cc_agent->GetKeyValue(absl::string_view(key, key_size),
absl::Seconds(timeout_seconds));
return ProcessGetKeyValueResult(value, status);
}
TF_Buffer* TF_CoordinationServiceTryGetKeyValue(
const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
TF_Status* status) {
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
auto value = cc_agent->TryGetKeyValue(absl::string_view(key, key_size));
return ProcessGetKeyValueResult(value, status);
}
void TF_CoordinationServiceDeleteKeyValue(const char* key, int64_t key_size,
TF_CoordinationServiceAgent* agent,
TF_Status* status) {
auto* cc_agent = reinterpret_cast<tsl::CoordinationServiceAgent*>(agent);
absl::Status cc_status =
cc_agent->DeleteKeyValue(absl::string_view(key, key_size));
status->status = cc_status;
}
// ---------------------------- PJRT -----------------------------------------
void TF_CreateAndSetPjRtCApiClient(const char* device_type, TF_Status* status,
PJRT_NamedValue* create_options,
int num_options) {
absl::StatusOr<std::unique_ptr<xla::PjRtClient>> pjrt_client =
xla::GetCApiClient(device_type, pjrt::ConvertFromPjRtNamedValueList(
create_options, num_options));
if (!pjrt_client.ok()) {
status->status = pjrt_client.status();
return;
}
absl::Status s = tensorflow::SetPjRtClientInTFGlobalResourceManager(
tensorflow::DeviceType(device_type), std::move(*pjrt_client));
status->status = s;
}
void TF_ResetPjRtCClient(const char* device_type, TF_Status* status) {
status->status =
tensorflow::ResetPjRtClient(tensorflow::DeviceType(device_type));
}
PJRT_Client* TF_GetPjRtCClient(const char* device_type, TF_Status* status) {
absl::StatusOr<xla::PjRtCApiClient*> pjrt_c_api_client =
tensorflow::GetPjRtCApiClient(tensorflow::DeviceType(device_type));
if (!pjrt_c_api_client.ok()) {
status->status = pjrt_c_api_client.status();
return nullptr;
}
status->status = absl::OkStatus();
return (*pjrt_c_api_client)->pjrt_c_client();
}
PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor, TF_Status* status) {
tensorflow::Tensor tensor;
auto s = tensorflow::TF_TensorToTensor(c_tensor, &tensor);
if (!s.ok()) {
status->status = s;
return nullptr;
}
absl::StatusOr<PJRT_Buffer*> c_buffer =
tensorflow::GetPjRtCBufferFromTensor(&tensor);
if (!c_buffer.ok()) {
status->status = c_buffer.status();
return nullptr;
}
status->status = absl::OkStatus();
return *c_buffer;
}
void TF_CreatePjRtBuffer(TF_Tensor* c_tensor, PJRT_Buffer* c_buffer,
const char* device_type, TF_Status* status) {
tensorflow::Tensor tensor;
auto s = tensorflow::TF_TensorToTensor(c_tensor, &tensor);
if (!s.ok()) {
status->status = s;
return;
}
absl::StatusOr<xla::PjRtCApiClient*> pjrt_c_api_client =
tensorflow::GetPjRtCApiClient(tensorflow::DeviceType(device_type));
if (!pjrt_c_api_client.ok()) {
status->status = pjrt_c_api_client.status();
return;
}
auto set_buffer_status =
SetPjRtCBufferToTensor(c_buffer, *pjrt_c_api_client, &tensor);
status->status = set_buffer_status;
}
tensorflow/c/experimental/next_pluggable_device/c_api.h
+156
View File
@@ -0,0 +1,156 @@
/* Copyright 2022 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_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_
#define TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_
#include <cstdint>
#include "tensorflow/c/c_api_macros.h"
#include "tensorflow/c/kernels.h"
#include "tensorflow/c/kernels_experimental.h"
#include "tensorflow/c/tf_buffer.h"
#include "tensorflow/c/tf_status.h"
#include "xla/pjrt/c/pjrt_c_api.h"
// --------------------------------------------------------------------------
// C API for device. The API is under active development and eventually
// should allow registering a plugin device with TensorFlow.
#ifdef __cplusplus
extern "C" {
#endif
// TF_Device is a C wrapper to the C++ TF Device class. This is to be passed
// through TF_OpKernelContext, and is opaque to plugin.
typedef struct TF_Device TF_Device;
typedef struct TF_VariableInfo TF_VariableInfo;
// Returns a `TF_Device` pointer, which actually points to a C++ `Device`.
// Currently we only allow `NextPluggableDevice` to be casted as `TF_Device`,
// but in theory every this is a C API for every kind of device.
TF_CAPI_EXPORT extern TF_Device* TF_GetDevice(TF_OpKernelContext* ctx);
// -------------------------- Resource ---------------------------------------
// Create a `tensorflow::PluginResource` to the ResourceMgr provided by the
// `ctx`. The `tensorflow::PluginResource` wraps a resource by plugin (as a
// opaque pointer, since TensorFlow cannot parse it). `delete_func` is needed
// for ResourceMgr to clean up the resource. `status` will be set.
TF_CAPI_EXPORT extern void TF_CreatePluginResource(
TF_OpKernelContext* ctx, const char* container_name,
const char* plugin_resource_name, void* plugin_resource,
void (*delete_func)(void*), TF_Status* status);
// If the ResourceMgr provided by the `ctx` has a resource
// `plugin_resource_name`, returns it in `*result_plugin_resource`. Otherwise,
// invokes create_func to create the resource. `delete_func` is needed for
// ResourceMgr to clean up the resource. `status` will be set. If `status` is
// not OK, `*result_plugin_resource` will be set as nullptr.
//
// Caller does not take ownership of the `plugin_resource`.
TF_CAPI_EXPORT extern void TF_LookupOrCreatePluginResource(
TF_OpKernelContext* ctx, const char* container_name,
const char* plugin_resource_name, void** result_plugin_resource,
void* (*create_func)(void*), void* create_func_args,
void (*delete_func)(void*), TF_Status* status);
// ------------------------- VariableInfo ------------------------------------
TF_CAPI_EXPORT extern TF_VariableInfo* TF_CreateVariableInfoFromContext(
TF_OpKernelContext* ctx, int index, TF_Status* status);
TF_CAPI_EXPORT extern void TF_LockVariableInfos(TF_VariableInfo** vars,
int num_vars,
TF_Status* status);
TF_CAPI_EXPORT extern void TF_AllocateTempForVariableInfo(
TF_OpKernelContext* ctx, TF_VariableInfo* var_info, TF_Status* status);
TF_CAPI_EXPORT extern TF_Tensor* TF_GetTensorFromVariableInfo(
TF_VariableInfo* var_info, TF_Status* status);
TF_CAPI_EXPORT extern void TF_DeleteVariableInfo(TF_VariableInfo* var_info);
// --------------------- Coordination service --------------------------------
// Returns a not owning pointer to the coordination service agent, which is
// opaque to plugin. Plugin OpKernels need to use the accompanying C APIs to
// access coordination service functionalities.
TF_CAPI_EXPORT extern TF_CoordinationServiceAgent*
TF_GetCoordinationServiceAgent(TF_OpKernelContext* ctx);
// Returns true if the coordination service agent has been initialized.
TF_CAPI_EXPORT extern bool TF_CoordinationServiceIsInitialized(
TF_CoordinationServiceAgent* agent);
TF_CAPI_EXPORT extern void TF_CoordinationServiceInsertKeyValue(
const char* key, int64_t key_size, const char* value, int64_t value_size,
TF_CoordinationServiceAgent* agent, TF_Status* status);
// Obtains key-value from coordination service agent. The returned `TF_Buffer`
// is a newly allocated buffer to hold the string key-value, and caller is
// responsible for managing the lifetime. If error, `status` will be set and a
// nullptr will be returned.
TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValue(
const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
TF_Status* status);
TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceGetKeyValueWithTimeout(
const char* key, int64_t key_size, int64_t timeout_seconds,
TF_CoordinationServiceAgent* agent, TF_Status* status);
TF_CAPI_EXPORT extern TF_Buffer* TF_CoordinationServiceTryGetKeyValue(
const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
TF_Status* status);
TF_CAPI_EXPORT extern void TF_CoordinationServiceDeleteKeyValue(
const char* key, int64_t key_size, TF_CoordinationServiceAgent* agent,
TF_Status* status);
// ---------------------------- PJRT -----------------------------------------
// Passes the pointer to a vector of PJRT_NamedValue and number of options to
// set options for creating a PJRT client. Passes nullptr for create_options and
// 0 for num_options if no options need to be set. You can use
// ConvertToPjRtNamedValueList in
// tensorflow/compiler/xla/pjrt/c/pjrt_c_api_helpers.h to generate the options.
TF_CAPI_EXPORT extern void TF_CreateAndSetPjRtCApiClient(
const char* device_type, TF_Status* status, PJRT_NamedValue* create_options,
int num_options);
// Resets the PjRt client for a device. After this, `TF_GetPjRtCClient` will
// returns an error for that device.
TF_CAPI_EXPORT extern void TF_ResetPjRtCClient(const char* device_type,
TF_Status* status);
// Gets the `PJRT_Client*` stored in TF global ResourceManager.
TF_CAPI_EXPORT extern PJRT_Client* TF_GetPjRtCClient(const char* device_type,
TF_Status* status);
// Gets the `PJRT_Buffer*` stored in the tensor. The status will contain error
// if the tensor does not have a `PjRtCApiBuffer`.
TF_CAPI_EXPORT extern PJRT_Buffer* TF_GetPjRtCBuffer(TF_Tensor* c_tensor,
TF_Status* status);
// Creates a `PjRtCApiBuffer` with the `PJRT_Buffer*` passed in and set to the
// tensor.
TF_CAPI_EXPORT extern void TF_CreatePjRtBuffer(TF_Tensor* c_tensor,
PJRT_Buffer* c_buffer,
const char* device_type,
TF_Status* status);
#ifdef __cplusplus
} // extern "C"
#endif
#endif // TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_C_API_H_
tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.cc
+92
View File
@@ -0,0 +1,92 @@
/* Copyright 2023 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/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
#include <memory>
#include <utility>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/jit/pjrt_tensor_buffer_util.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/tsl/platform/errors.h"
#include "xla/tsl/platform/statusor.h"
#include "tensorflow/core/framework/resource_mgr.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/global_state.h"
#include "tensorflow/core/tfrt/common/pjrt_state.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
namespace tensorflow {
absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor) {
tensorflow::AsyncValueTensor* av_tensor =
tensorflow::AsyncValueTensor::FromTensor(tensor);
if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {
return absl::InternalError("Input tensor does not have PjRtBuffer.");
}
auto* c_api_buffer =
dynamic_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
if (c_api_buffer == nullptr) {
return absl::InternalError(
"The PjRtBuffer in the tensor is not type PjRtCApiBuffer.");
}
return c_api_buffer->c_buffer();
}
absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
xla::PjRtCApiClient* c_api_client,
Tensor* tensor) {
auto buffer = std::make_unique<xla::PjRtCApiBuffer>(c_api_client, c_buffer);
tensorflow::AsyncValueTensor* av_tensor =
tensorflow::AsyncValueTensor::FromTensor(tensor);
if (av_tensor == nullptr) {
TF_ASSIGN_OR_RETURN(
*tensor, MakeTensorFromPjRtBuffer(tensor->dtype(), tensor->shape(),
std::move(buffer)));
} else {
av_tensor->SetBuffer(std::move(buffer));
}
return absl::OkStatus();
}
absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
const DeviceType& device_type) {
TF_ASSIGN_OR_RETURN(absl::StatusOr<xla::PjRtClient*> pjrt_client,
tensorflow::GetPjRtClient(device_type));
auto* pjrt_c_api_client = dynamic_cast<xla::PjRtCApiClient*>(*pjrt_client);
if (pjrt_c_api_client == nullptr) {
return absl::InternalError(absl::StrCat("PjRtClient for ",
device_type.type_string(),
" is not type PjRtCApiClient"));
}
return pjrt_c_api_client;
}
absl::Status ResetPjRtClient(const DeviceType& device_type) {
ResourceMgr* rmgr = tfrt_global::GetTFGlobalResourceMgr();
PjRtState* pjrt_state;
TF_RETURN_IF_ERROR(rmgr->Lookup(rmgr->default_container(),
kPjRtStateResourceName, &pjrt_state));
TF_RETURN_IF_ERROR(pjrt_state->MovePjRtClientToUnused(device_type));
return absl::OkStatus();
}
} // namespace tensorflow
tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h
+39
View File
@@ -0,0 +1,39 @@
/* Copyright 2023 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_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_
#define TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
#include "tensorflow/core/framework/tensor.h"
namespace tensorflow {
absl::StatusOr<PJRT_Buffer*> GetPjRtCBufferFromTensor(const Tensor* tensor);
absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
xla::PjRtCApiClient* c_api_client,
Tensor* tensor);
absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
const DeviceType& device_type);
absl::Status ResetPjRtClient(const DeviceType& device_type);
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_NEXT_PLUGGABLE_DEVICE_TENSOR_PJRT_BUFFER_UTIL_H_
tensorflow/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util_test.cc
+206
View File
@@ -0,0 +1,206 @@
/* Copyright 2023 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/c/experimental/next_pluggable_device/tensor_pjrt_buffer_util.h"
#include <cstdint>
#include <memory>
#include <optional>
#include <utility>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/base/casts.h"
#include "absl/log/check.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_cat.h"
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c/pjrt_c_api_cpu.h"
#include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
#include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
#include "xla/pjrt/pjrt_api.h"
#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
#include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/status_matchers.h"
#include "xla/tsl/platform/statusor.h"
#include "xla/tsl/protobuf/error_codes.pb.h"
#include "xla/xla_data.pb.h"
#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/tfrt/common/async_value_tensor.h"
#include "tensorflow/core/tfrt/common/pjrt_util.h"
#include "tsl/platform/casts.h"
namespace tensorflow {
namespace {
using ::testing::HasSubstr;
using ::testing::NotNull;
using ::tsl::testing::StatusIs;
PJRT_Buffer* CreateCBuffer() {
auto status = pjrt::PjrtApi(DEVICE_CPU);
if (!status.ok()) {
CHECK_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
}
auto pjrt_client = xla::GetCApiClient(DEVICE_CPU);
CHECK_OK(pjrt_client.status());
auto c_api_client = absl::down_cast<xla::PjRtCApiClient*>(pjrt_client->get());
std::vector<int32_t> data(1, 0);
xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
auto buffer = c_api_client->pjrt_c_client()->client->BufferFromHostBuffer(
data.data(), shape.element_type(), shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr,
c_api_client->pjrt_c_client()->client->memory_spaces()[0],
/*device_layout=*/nullptr);
CHECK_OK(buffer.status());
return new PJRT_Buffer{std::move(*buffer), c_api_client->pjrt_c_client()};
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorNoBuffer) {
auto allocator = std::make_unique<AsyncValueAllocator>();
tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
EXPECT_THAT(
GetPjRtCBufferFromTensor(&tensor),
absl_testing::StatusIs(
error::INTERNAL,
HasSubstr(absl::StrCat("Input tensor does not have PjRtBuffer"))));
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorIncoorectType) {
auto allocator = std::make_unique<AsyncValueAllocator>();
tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
xla::CpuClientOptions options;
options.asynchronous = true;
options.cpu_device_count = 1;
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
std::vector<int32_t> data(1, 0);
xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
TF_ASSERT_OK_AND_ASSIGN(
auto buffer,
pjrt_client->BufferFromHostBuffer(
data.data(), shape.element_type(), shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
nullptr, pjrt_client->memory_spaces()[0], /*device_layout=*/nullptr));
tensorflow::AsyncValueTensor* av_tensor =
tensorflow::AsyncValueTensor::FromTensor(&tensor);
av_tensor->SetBuffer(std::move(buffer));
EXPECT_THAT(
GetPjRtCBufferFromTensor(&tensor),
absl_testing::StatusIs(
error::INTERNAL,
HasSubstr(absl::StrCat(
"The PjRtBuffer in the tensor is not type PjRtCApiBuffer"))));
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorSuccess) {
auto allocator = std::make_unique<AsyncValueAllocator>();
tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
auto status = pjrt::PjrtApi(DEVICE_CPU);
if (!status.ok()) {
TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
}
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
std::vector<int32_t> data(1, 0);
xla::Shape shape = xla::ShapeUtil::MakeShape(xla::S32, {1});
TF_ASSERT_OK_AND_ASSIGN(
auto buffer,
pjrt_client->BufferFromHostBuffer(
data.data(), shape.element_type(), shape.dimensions(),
/*byte_strides=*/std::nullopt,
xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
nullptr, pjrt_client->memory_spaces()[0], /*device_layout=*/nullptr));
tensorflow::AsyncValueTensor* av_tensor =
tensorflow::AsyncValueTensor::FromTensor(&tensor);
av_tensor->SetBuffer(std::move(buffer));
TF_ASSERT_OK_AND_ASSIGN(auto c_buffer, GetPjRtCBufferFromTensor(&tensor));
EXPECT_THAT(c_buffer, NotNull());
}
TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorNotAsyncValueTensor) {
tensorflow::Tensor tensor(DT_FLOAT, {1});
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
PJRT_Buffer* c_buffer = CreateCBuffer();
TF_EXPECT_OK(SetPjRtCBufferToTensor(
c_buffer, absl::down_cast<xla::PjRtCApiClient*>(pjrt_client.get()),
&tensor));
}
TEST(TensorPjRtBufferUtilTest, SetPjRtCBufferToTensorSuccess) {
auto allocator = std::make_unique<AsyncValueAllocator>();
tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
PJRT_Buffer* c_buffer = CreateCBuffer();
TF_EXPECT_OK(SetPjRtCBufferToTensor(
c_buffer, absl::down_cast<xla::PjRtCApiClient*>(pjrt_client.get()),
&tensor));
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientNotFound) {
EXPECT_THAT(GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)),
absl_testing::StatusIs(
error::NOT_FOUND,
HasSubstr(absl::StrCat(
"PjRt client not found for device type ", DEVICE_CPU))));
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientIncorrectType) {
xla::CpuClientOptions options;
options.asynchronous = true;
options.cpu_device_count = 1;
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetXlaPjrtCpuClient(options));
TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU,
std::move(pjrt_client)));
EXPECT_THAT(GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)),
absl_testing::StatusIs(
error::INTERNAL,
HasSubstr(absl::StrCat("PjRtClient for ", DEVICE_CPU,
" is not type PjRtCApiClient"))));
}
TEST(TensorPjRtBufferUtilTest, GetPjRtCApiClientSuccess) {
auto status = pjrt::PjrtApi(DEVICE_CPU);
if (!status.ok()) {
TF_ASSERT_OK(pjrt::SetPjrtApi(DEVICE_CPU, GetPjrtApi()));
}
TF_ASSERT_OK_AND_ASSIGN(auto pjrt_client, xla::GetCApiClient(DEVICE_CPU));
TF_ASSERT_OK(SetPjRtClientInTFGlobalResourceManager(DEVICE_CPU,
std::move(pjrt_client)));
TF_ASSERT_OK_AND_ASSIGN(
auto pjrt_client_get,
GetPjRtCApiClient(tensorflow::DeviceType(DEVICE_CPU)));
EXPECT_THAT(pjrt_client_get, NotNull());
}
} // namespace
} // namespace tensorflow
tensorflow/c/experimental/ops/BUILD
+165
View File
@@ -0,0 +1,165 @@
load("//tensorflow:tensorflow.default.bzl", "filegroup")
load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
# Experimental ops. These will eventually be replaced by machine-generated versions.
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
cc_library(
name = "array_ops",
srcs = [
"array_ops.cc",
],
hdrs = [
"array_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:tracing_utils",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "io_ops",
srcs = [
"io_ops.cc",
],
hdrs = [
"io_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:tracing_utils",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "math_ops",
srcs = [
"math_ops.cc",
],
hdrs = [
"math_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:tracing_utils",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "nn_ops",
srcs = [
"nn_ops.cc",
],
hdrs = [
"nn_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:tracing_utils",
"//tensorflow/core:framework_types_hdr",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "resource_variable_ops",
srcs = [
"resource_variable_ops.cc",
],
hdrs = [
"resource_variable_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_operation",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/c/eager:tracing_utils",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:errors",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
cc_library(
name = "ops",
hdrs = [
"array_ops.h",
"io_ops.h",
"math_ops.h",
"nn_ops.h",
"resource_variable_ops.h",
],
visibility = [
"//tensorflow:internal",
],
deps = [
":array_ops",
":io_ops",
":math_ops",
":nn_ops",
":resource_variable_ops",
"//tensorflow/c/eager:abstract_context",
"//tensorflow/c/eager:abstract_tensor_handle",
"//tensorflow/core:framework",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:span",
],
)
filegroup(
name = "pywrap_required_hdrs",
srcs = [
"array_ops.h",
"io_ops.h",
"math_ops.h",
"nn_ops.h",
"resource_variable_ops.h",
],
visibility = ["//tensorflow/python:__subpackages__"],
)
tensorflow/c/experimental/ops/README.md
+20
View File
@@ -0,0 +1,20 @@
# Experimental C++ Ops
The C++ files in this directory (***\*_ops.h*** and ***\*_ops.cc***) are
autogenerated from the registered Op and API definitions.
To regenerate them, run the script in this directory, `update_cpp_ops.sh`, with
no arguments. This script will overwrite the existing ops in-place at
***tensorflow/c/experimental/ops/\*_ops.{cc,h}***.
Run this `update_cpp_ops.sh` script when Op definitions change in the registry.
To generate additional operators, extend the lists in this script. Note that
category names correspond to generated source file names, and should be
consistent with the original source files registering each operator. For example
since `REGISTER_OP("MatMul")` appears in ***core/math_ops.cc***, the "MatMul"
operator in the script should be in the "math" category, and it will be
generated in the output file `c/experimental/ops/math_ops.cc`.
Running this script should be a no-op, generating identical code other than
formatting (i.e., line wrapping).
tensorflow/c/experimental/ops/array_ops.cc
+210
View File
@@ -0,0 +1,210 @@
/* Copyright 2021 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.
==============================================================================*/
// This file is MACHINE GENERATED! Do not edit.
#include "tensorflow/c/experimental/ops/array_ops.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_operation.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/tracing_utils.h"
#include "tensorflow/core/framework/types.h" // NOLINT
#include "tensorflow/core/framework/types.pb.h"
#include "tensorflow/core/platform/errors.h" // NOLINT
using tensorflow::tracing::MaybeSetOpName;
namespace tensorflow {
namespace ops {
// Op: Identity()
// Summary: Return a tensor with the same shape and contents as the input tensor
// or value.
//
// Description:
absl::Status Identity(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle** output, const char* name,
const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("Identity", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(input));
int num_retvals = 1;
TF_RETURN_IF_ERROR(op_ptr->Execute(absl::MakeSpan(output, 1), &num_retvals));
if (num_retvals != 1) {
return absl::InternalError("Identity: unexpected number of outputs");
}
return absl::OkStatus();
}
// Op: IdentityN()
// Summary: Returns a list of tensors with the same shapes and contents as the
// input
//
// Description:
// tensors.
//
// This op can be used to override the gradient for complicated functions. For
// example, suppose y = f(x) and we wish to apply a custom function g for
// backprop such that dx = g(dy). In Python,
//
// ```python
// with tf.get_default_graph().gradient_override_map(
// {'IdentityN': 'OverrideGradientWithG'}):
// y, _ = identity_n([f(x), x])
//
// @tf.RegisterGradient('OverrideGradientWithG')
// def ApplyG(op, dy, _):
// return [None, g(dy)] # Do not backprop to f(x).
// ```
absl::Status IdentityN(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> input,
absl::Span<AbstractTensorHandle*> output,
const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("IdentityN", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInputList(input));
int num_retvals = output.size();
TF_RETURN_IF_ERROR(op_ptr->Execute(output, &num_retvals));
if (num_retvals != output.size()) {
return absl::InternalError("IdentityN: unexpected number of outputs");
}
return absl::OkStatus();
}
// Op: ZerosLike()
// Summary: Returns a tensor of zeros with the same shape and type as x.
//
// Description:
absl::Status ZerosLike(AbstractContext* ctx, AbstractTensorHandle* const x,
AbstractTensorHandle** y, const char* name,
const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("ZerosLike", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(x));
int num_retvals = 1;
TF_RETURN_IF_ERROR(op_ptr->Execute(absl::MakeSpan(y, 1), &num_retvals));
if (num_retvals != 1) {
return absl::InternalError("ZerosLike: unexpected number of outputs");
}
return absl::OkStatus();
}
// Op: Shape()
// Summary: Returns the shape of a tensor.
//
// Description:
// This operation returns a 1-D integer tensor representing the shape of
// `input`.
//
// For example:
//
// ```
// # 't' is [[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]]
// shape(t) ==> [2, 2, 3]
// ```
absl::Status Shape(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle** output, DataType out_type,
const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("Shape", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(input));
TF_RETURN_IF_ERROR(op_ptr->SetAttrType("out_type", out_type));
int num_retvals = 1;
TF_RETURN_IF_ERROR(op_ptr->Execute(absl::MakeSpan(output, 1), &num_retvals));
if (num_retvals != 1) {
return absl::InternalError("Shape: unexpected number of outputs");
}
return absl::OkStatus();
}
// Op: ExpandDims()
// Summary: Inserts a dimension of 1 into a tensor's shape.
//
// Description:
// Given a tensor `input`, this operation inserts a dimension of 1 at the
// dimension index `axis` of `input`'s shape. The dimension index `axis`
// starts at zero; if you specify a negative number for `axis` it is counted
// backward from the end.
//
// This operation is useful if you want to add a batch dimension to a single
// element. For example, if you have a single image of shape `[height, width,
// channels]`, you can make it a batch of 1 image with `expand_dims(image,
// 0)`, which will make the shape `[1, height, width, channels]`.
//
// Other examples:
//
// ```
// # 't' is a tensor of shape [2]
// shape(expand_dims(t, 0)) ==> [1, 2]
// shape(expand_dims(t, 1)) ==> [2, 1]
// shape(expand_dims(t, -1)) ==> [2, 1]
//
// # 't2' is a tensor of shape [2, 3, 5]
// shape(expand_dims(t2, 0)) ==> [1, 2, 3, 5]
// shape(expand_dims(t2, 2)) ==> [2, 3, 1, 5]
// shape(expand_dims(t2, 3)) ==> [2, 3, 5, 1]
// ```
//
// This operation requires that:
//
// `-1-input.dims() <= dim <= input.dims()`
//
// This operation is related to `squeeze()`, which removes dimensions of
// size 1.
absl::Status ExpandDims(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle* const dim,
AbstractTensorHandle** output, const char* name,
const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("ExpandDims", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(input));
TF_RETURN_IF_ERROR(op_ptr->AddInput(dim));
int num_retvals = 1;
TF_RETURN_IF_ERROR(op_ptr->Execute(absl::MakeSpan(output, 1), &num_retvals));
if (num_retvals != 1) {
return absl::InternalError("ExpandDims: unexpected number of outputs");
}
return absl::OkStatus();
}
// Op: OnesLike()
// Summary: Returns a tensor of ones with the same shape and type as x.
//
// Description:
absl::Status OnesLike(AbstractContext* ctx, AbstractTensorHandle* const x,
AbstractTensorHandle** y, const char* name,
const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("OnesLike", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(x));
int num_retvals = 1;
TF_RETURN_IF_ERROR(op_ptr->Execute(absl::MakeSpan(y, 1), &num_retvals));
if (num_retvals != 1) {
return absl::InternalError("OnesLike: unexpected number of outputs");
}
return absl::OkStatus();
}
} // namespace ops
} // namespace tensorflow
tensorflow/c/experimental/ops/array_ops.h
+70
View File
@@ -0,0 +1,70 @@
/* Copyright 2021 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.
==============================================================================*/
// This file is MACHINE GENERATED! Do not edit.
#ifndef TENSORFLOW_C_EXPERIMENTAL_OPS_ARRAY_OPS_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_ARRAY_OPS_H_
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/core/framework/types.h" // NOLINT
#include "tensorflow/core/framework/types.pb.h"
namespace tensorflow {
namespace ops {
// Return a tensor with the same shape and contents as the input tensor or
// value.
absl::Status Identity(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle** output, const char* name = nullptr,
const char* raw_device_name = nullptr);
// Returns a list of tensors with the same shapes and contents as the input
absl::Status IdentityN(AbstractContext* ctx,
absl::Span<AbstractTensorHandle* const> input,
absl::Span<AbstractTensorHandle*> output,
const char* name = nullptr,
const char* raw_device_name = nullptr);
// Returns a tensor of zeros with the same shape and type as x.
absl::Status ZerosLike(AbstractContext* ctx, AbstractTensorHandle* const x,
AbstractTensorHandle** y, const char* name = nullptr,
const char* raw_device_name = nullptr);
// Returns the shape of a tensor.
absl::Status Shape(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle** output, DataType out_type = DT_INT32,
const char* name = nullptr,
const char* raw_device_name = nullptr);
// Inserts a dimension of 1 into a tensor's shape.
absl::Status ExpandDims(AbstractContext* ctx, AbstractTensorHandle* const input,
AbstractTensorHandle* const dim,
AbstractTensorHandle** output,
const char* name = nullptr,
const char* raw_device_name = nullptr);
// Returns a tensor of ones with the same shape and type as x.
absl::Status OnesLike(AbstractContext* ctx, AbstractTensorHandle* const x,
AbstractTensorHandle** y, const char* name = nullptr,
const char* raw_device_name = nullptr);
} // namespace ops
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_ARRAY_OPS_H_
tensorflow/c/experimental/ops/gen/BUILD
+36
View File
@@ -0,0 +1,36 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_binary",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "generate_cpp_main",
srcs = ["generate_cpp_main.cc"],
deps = [
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings",
"//tensorflow/c/experimental/ops/gen/common",
"//tensorflow/c/experimental/ops/gen/cpp/renderers",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/c/experimental/ops/gen/cpp",
"//tensorflow/c/experimental/ops/gen/model",
"//tensorflow/core:lib",
# Without this line to link in ops, the global registry will be empty:
"//tensorflow/core:ops",
],
)
tf_cc_binary(
name = "generate_cpp",
args = ["--stderrthreshold=1"], # Warnings and errors to stderr.
deps = [":generate_cpp_main"],
)
tensorflow/c/experimental/ops/gen/README.md
+116
View File
@@ -0,0 +1,116 @@
# TensorFlow Op CodeGen Machinery (Experimental)
## Usage
```
usage: generate_cpp [flags] OpName1 [OpName2 ...]
Flags:
--help=false bool Print this help message.
--category="" string Category for generated ops (e.g. 'math', 'array').
--namespace="" string Compact C++ namespace, default is 'tensorflow::ops'.
--output_dir="" string Directory into which output files will be generated.
--source_dir="" string The tensorflow root directory, e.g. 'tensorflow/' for in-source include paths. Any path underneath the tensorflow root is also accepted.
--api_dirs="" string Comma-separated list of directories containing API definitions.
```
## Design
### Generator Framework
The generator framework is a loose Model/View/Controller arrangement:
The *Model* classes live in the ***model/*** directory. They are representations
of the `OpDef` and `ApiDef` protos, normalized and resolved.
> _For example, an `OpDef` proto's `ArgDef` members contain a type string, which
> must be dereferenced to an `AttrDef` by name to determine its type. This
> `AttrDef` proto message in turn contains a type string which may need to be
> parsed as "list(type)". Other `AttrDef` messages are not types, but instead
> argument-like modifiers. In contrast, the generator model `ArgSpec` contains a
> resolved `ArgType` which provides a boolean `is_list()` method directly, and
> the model `OpSpec` provides a list of only the argument-like attributes. In
> addition to convenience, this should aid consistency between generated code in
> each target language._
The *Controller* is in the ***common/*** directory. It is the workhorse used by
the language generators; it digests the Op registry and API definitions to build
the model and provides utilities for the language generators.
The *View* and rendering classes map the language-independent Model classes
(`OpSpec`, `ArgSpec`, `AttrSpec`, etc.) to language-specific `SourceCode`. The
framework does not impose any design on the language-specific generators, but
provides some utilities, and the C++ generator is a complete example.
### C++ Generator
The `CppGenerator` class is the interface to the `cpp/` language directory.
Given a config, it can generate source code for a .cc or .h file as a string or
write it to a target file.
The `CppFileRenderer` is the main renderer used by the generator; it renders an
entire file. The `CppConfig` defines if it is operating in header or source
mode.
"Views" are stateless and intended to be low-level building blocks: a direct
language-specific representation of the model classes. For example, an `ArgView`
is initialized from an `ArgSpec` (which was created initially from an `ArgDef`
proto message). Where they may have some similar methods between the model and
view, the view methods are language-specific.
For instance, the C++ generator's `ArgView::VariableName()` method is an
language-formatted name usable as a variable representing the model `ArgSpec`
object. In contrast, the `ArgSpec::name()` method in the model refers to the
canonical name of the object in the proto.
Where views are a representation of the *input* model, in the C++ generator,
"renderers" then use these views to build the *output* `SourceCode`; Renderers
understand the language at the statement/directive level and target a functional
section of the output, such as a block comment or an entire method or file.
Other differences between views and renderers:
* Renderers are stateful, modifying a referenced SourceCode. Views are
stateless and their public methods are all const, returning strings.
* Renderers are context-dependent, e.g. a method signature will include
default values when in "declaration" mode but not "definition" mode. A view
of some argument object simply knows its default value and does not care the
context.
* In terms of dependencies, `Renderers` use `Views` and other `Renderers`.
However, `Renderers` do **not** reference the model directly (e.g.
`OpSpec`). This is because if a renderer needs to reference part of the
model, it should get a language specific representation.
### Extending to Additional Languages
The design for the C++ generator should apply to other languages, and the
underlying generator framework (the model and controller) try to be agnostic. In
fact, some elements of the C++ design could be formalized (such as the
rendering/view framework) or re-used (e.g. `cpp:Renderer` could likely be shared
with C and Java as a common C-style language renderer base class).
Abstracted and condensed from the C++ generator, the overall control flow could
be described as follows:
From main() in *generate_lang_main.cc*:
* Call `tensorflow::port::InitMain` and parse any flags
* Initialize config objects (e.g. `PathConfig`, `LangConfig` from flags)
* Initialize a new `LangGenerator` from these config objects
* Call this generator to create/write `SourceCode` to a file
In class `LangGenerator` in *lang_generator.cc*:
* Initialize a new `Controller` from the config objects
* Call this controller to build the Op models (`OpSpec`)
* Initialize a new language-specific `View` for each model object
* Create a blank `SourceCode` rendering target (for each output file)
* Initialize a new `LangFileRenderer` from this target source code, the model
`View` objects, and config objects
* Call this renderer to generate the target `SourceCode`
The dependencies are as follows:
* `lang::Generator` depends on `Controller`, `Model`, `lang::Renderers`,
`lang::Views`
* `lang::Renderer` depends on `lang::View` (and `lang::Renderer` peers)
* `lang::View` depends on the model (e.g. `OpSpec`) (and `lang::View` peers)
tensorflow/c/experimental/ops/gen/common/BUILD
+49
View File
@@ -0,0 +1,49 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_tests",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
cc_library(
name = "common",
srcs = glob(
["*.cc"],
exclude = ["*_test.cc"],
),
hdrs = glob(["*.h"]),
deps = [
"//tensorflow/c/experimental/ops/gen/model",
"//tensorflow/core:framework",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:op_gen_lib",
"//tensorflow/core:protos_all_cc",
"//tensorflow/core/platform:str_util",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
tf_cc_tests(
name = "all_tests",
size = "small",
srcs = glob(["*_test.cc"]),
deps = [
":common",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/platform:types",
],
)
tensorflow/c/experimental/ops/gen/common/case_format.cc
+107
View File
@@ -0,0 +1,107 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/case_format.h"
#include <string>
#include "absl/strings/ascii.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace {
enum CaseFormatType {
LOWER_CAMEL,
UPPER_CAMEL,
LOWER_SNAKE,
UPPER_SNAKE,
};
std::string FormatStringCase(const std::string& str, CaseFormatType to,
const char delimiter = '_') {
const bool from_snake = (str == absl::AsciiStrToUpper(str)) ||
(str == absl::AsciiStrToLower(str));
const bool toUpper = (to == UPPER_CAMEL || to == UPPER_SNAKE);
const bool toSnake = (to == LOWER_SNAKE || to == UPPER_SNAKE);
std::string result;
bool inputStart = true;
bool wordStart = true;
for (const char c : str) {
// Find a word start.
if (c == delimiter) {
// Repeated cases of wordStart means explicit delimiter usage.
if (wordStart) {
result.push_back(delimiter);
}
wordStart = true;
continue;
}
if (!from_snake && absl::ascii_isupper(c)) {
wordStart = true;
}
// add delimiter
if (wordStart && toSnake && !inputStart) {
result.push_back(delimiter);
}
// add the next letter from the input string (choosing upper/lower case)
const bool shouldCapIfSnake = toUpper;
const bool shouldCapIfCamel = wordStart && (toUpper || !inputStart);
if ((toSnake && shouldCapIfSnake) || (!toSnake && shouldCapIfCamel)) {
result += absl::ascii_toupper(c);
} else {
result += absl::ascii_tolower(c);
}
// at this point we are no longer at the start of a word:
wordStart = false;
// .. or the input:
inputStart = false;
}
if (wordStart) {
// This only happens with a trailing delimiter, which should remain.
result.push_back(delimiter);
}
return result;
}
} // namespace
//
// Public interface
//
std::string toLowerCamel(const std::string& s, const char delimiter) {
return FormatStringCase(s, LOWER_CAMEL, delimiter);
}
std::string toLowerSnake(const std::string& s, const char delimiter) {
return FormatStringCase(s, LOWER_SNAKE, delimiter);
}
std::string toUpperCamel(const std::string& s, const char delimiter) {
return FormatStringCase(s, UPPER_CAMEL, delimiter);
}
std::string toUpperSnake(const std::string& s, const char delimiter) {
return FormatStringCase(s, UPPER_SNAKE, delimiter);
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/case_format.h
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_COMMON_CASE_FORMAT_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_CASE_FORMAT_H_
#include "tensorflow/core/platform/str_util.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
// Conversion routines between upper/lower camel/snake case formats, e.g.:
// "lowerCamelCase"
// "lower_snake_case"
// "UpperCamelCase"
// "UPPER_SNAKE_CASE"
//
// The input format is automatically detected.
// The delimiter must be specified if it is other than an underscore ('_')
// for conversion either *to* or *from* snake case.
//
// Leading and trailing delimiters are supported, e.g.:
// "__OneTwo__" (in camel case) <==> "__ONE_TWO__" (in snake case)
//
// Note: performance not yet tested.
std::string toLowerCamel(const std::string& s, const char delimiter = '_');
std::string toLowerSnake(const std::string& s, const char delimiter = '_');
std::string toUpperCamel(const std::string& s, const char delimiter = '_');
std::string toUpperSnake(const std::string& s, const char delimiter = '_');
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_CASE_FORMAT_H_
tensorflow/c/experimental/ops/gen/common/case_format_test.cc
+128
View File
@@ -0,0 +1,128 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/case_format.h"
#include <string>
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace {
// For each test case, we manually construct the 4 variations in string case and
// test all 16 conversions: from and to each of the 4 string case variations.
struct Variations {
std::string lower_camel;
std::string lower_snake;
std::string upper_camel;
std::string upper_snake;
};
void TestSingleVariation(const std::string& str, Variations expected,
char delimiter = '_') {
EXPECT_EQ(expected.lower_camel, toLowerCamel(str, delimiter));
EXPECT_EQ(expected.lower_snake, toLowerSnake(str, delimiter));
EXPECT_EQ(expected.upper_camel, toUpperCamel(str, delimiter));
EXPECT_EQ(expected.upper_snake, toUpperSnake(str, delimiter));
}
void TestAllVariations(Variations variations, char delimiter = '_') {
TestSingleVariation(variations.lower_camel, variations, delimiter);
TestSingleVariation(variations.lower_snake, variations, delimiter);
TestSingleVariation(variations.upper_camel, variations, delimiter);
TestSingleVariation(variations.upper_snake, variations, delimiter);
}
TEST(CppOpGenCaseFormat, test_single_word) {
TestAllVariations(Variations{
"three",
"three",
"Three",
"THREE",
});
}
TEST(CppOpGenCaseFormat, test_complex_string) {
TestAllVariations(Variations{
"threeNTest33Words",
"three_n_test33_words",
"ThreeNTest33Words",
"THREE_N_TEST33_WORDS",
});
}
TEST(CppOpGenCaseFormat, test_hyphen_delimiter) {
TestAllVariations(
Variations{
"threeNTest33Words",
"three-n-test33-words",
"ThreeNTest33Words",
"THREE-N-TEST33-WORDS",
},
'-');
}
TEST(CppOpGenCaseFormat, test_trailing_underscore) {
TestAllVariations(Variations{
"threeNTest33Words_",
"three_n_test33_words_",
"ThreeNTest33Words_",
"THREE_N_TEST33_WORDS_",
});
}
TEST(CppOpGenCaseFormat, test_double_trailing_underscores) {
TestAllVariations(Variations{
"xxY__",
"xx_y__",
"XxY__",
"XX_Y__",
});
}
TEST(CppOpGenCaseFormat, test_leading_underscore) {
TestAllVariations(Variations{
"_threeNTest33Words",
"_three_n_test33_words",
"_ThreeNTest33Words",
"_THREE_N_TEST33_WORDS",
});
}
TEST(CppOpGenCaseFormat, test_double_leading_underscores) {
TestAllVariations(Variations{
"__threeNTest33Words",
"__three_n_test33_words",
"__ThreeNTest33Words",
"__THREE_N_TEST33_WORDS",
});
}
TEST(CppOpGenCaseFormat, test_leading_and_trailing_underscores) {
TestAllVariations(Variations{
"__threeNTest33Words____",
"__three_n_test33_words____",
"__ThreeNTest33Words____",
"__THREE_N_TEST33_WORDS____",
});
}
} // namespace
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/controller.cc
+94
View File
@@ -0,0 +1,94 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/controller.h"
#include <string>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/substitute.h"
#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
#include "tensorflow/c/experimental/ops/gen/common/source_code.h"
#include "tensorflow/c/experimental/ops/gen/model/op_spec.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/framework/api_def.pb.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/op_gen_lib.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/path.h"
namespace tensorflow {
namespace generator {
Controller::Controller(PathConfig path_config, Env* env)
: env_(env), path_config_(path_config) {
// Load the Op and API definitions
InitializeOpApi();
// Convert the Op and API definitions to the internal data model
BuildModel();
}
Controller::~Controller() { delete api_def_map_; }
const void Controller::WriteFile(const std::string& file_path,
const SourceCode& code) const {
TF_CHECK_OK(WriteStringToFile(env_, file_path, code.Render())) << file_path;
}
const std::vector<OpSpec>& Controller::GetModelOps() const {
return operators_;
}
void Controller::InitializeOpApi() {
OpRegistry::Global()->Export(false, &op_list_);
// Load matching API defs for each Op. Paths are visited in order, allowing
// python/api_def_Xyz.pbtxt to override base/api_def_Xyz.pbtxt, for example.
api_def_map_ = new ApiDefMap(op_list_);
for (const auto& op : op_list_.op()) {
for (const auto& dir : path_config_.api_dirs) {
const std::string file_name =
absl::Substitute("api_def_$0.pbtxt", op.name());
const std::string file_path = io::JoinPath(dir, file_name);
if (env_->FileExists(file_path).ok()) {
TF_CHECK_OK(api_def_map_->LoadFile(env_, file_path)) << file_path;
} else {
// API defs are currently used for only optional pieces.
}
}
}
// Doc strings (summary, description) typically come from the API def.
api_def_map_->UpdateDocs();
}
void Controller::BuildModel() {
// Build the internal data model for the requested ops
for (const auto& op_name : path_config_.op_names) {
const OpDef* op_def = nullptr;
TF_CHECK_OK(OpRegistry::Global()->LookUpOpDef(op_name, &op_def));
CHECK(op_def != nullptr); // Crash OK
const ApiDef* api_def = api_def_map_->GetApiDef(op_name);
CHECK(api_def != nullptr); // Crash OK
operators_.push_back(OpSpec::Create(*op_def, *api_def));
}
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/controller.h
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_COMMON_CONTROLLER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_CONTROLLER_H_
#include <vector>
#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
#include "tensorflow/c/experimental/ops/gen/common/source_code.h"
#include "tensorflow/c/experimental/ops/gen/model/op_spec.h"
#include "tensorflow/core/framework/op_def.pb.h"
#include "tensorflow/core/framework/op_gen_lib.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
class Controller {
public:
explicit Controller(PathConfig path_config, Env* env = Env::Default());
virtual ~Controller();
const void WriteFile(const std::string& file_path,
const SourceCode& code) const;
const std::vector<OpSpec>& GetModelOps() const;
private:
void InitializeOpApi();
void BuildModel();
// Data model: Ops to generate
std::vector<OpSpec> operators_;
// Configuration
Env* env_;
PathConfig path_config_;
// Initialized TensorFlow Op/API definitions
OpList op_list_;
ApiDefMap* api_def_map_;
};
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_CONTROLLER_H_
tensorflow/c/experimental/ops/gen/common/path_config.cc
+69
View File
@@ -0,0 +1,69 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/path_config.h"
#include <algorithm>
#include <string>
#include <vector>
#include "absl/strings/str_join.h"
#include "tensorflow/core/lib/strings/str_util.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
PathConfig::PathConfig(const std::string& output_dir,
const std::string& source_dir,
const std::string& api_dir_list,
const std::vector<std::string> op_names)
: output_path(output_dir), op_names(op_names) {
api_dirs = str_util::Split(api_dir_list, ",", str_util::SkipEmpty());
// Decompose the directory components given the output/source directories.
//
// Be flexible; accept any path string with a "tensorflow" directory name.
// (It's hard to construct a location-agnostic include path string using the
// build system, so we accept an example, such as the source build target.)
tf_root_dir = "tensorflow";
// Prefix, e.g. "third_party" given root_dir "third_party/tensorflow/...."
std::vector<std::string> source_path_components =
tensorflow::str_util::Split(source_dir, "/");
auto source_tfroot_pos = std::find(source_path_components.begin(),
source_path_components.end(), tf_root_dir);
if (source_tfroot_pos != source_path_components.end()) {
tf_prefix_dir =
absl::StrJoin(source_path_components.begin(), source_tfroot_pos, "/");
} else {
tf_prefix_dir = source_dir;
}
// TF subdir, e.g. "c/ops" given output_dir "blah/blah/tensorflow/c/ops"
std::vector<std::string> output_path_components =
tensorflow::str_util::Split(output_dir, "/");
auto output_tfroot_pos = std::find(output_path_components.begin(),
output_path_components.end(), tf_root_dir);
if (output_tfroot_pos != output_path_components.end()) {
tf_output_dir =
absl::StrJoin(output_tfroot_pos + 1, output_path_components.end(), "/");
} else {
tf_output_dir = output_dir;
}
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/path_config.h
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_COMMON_PATH_CONFIG_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_PATH_CONFIG_H_
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
struct PathConfig {
std::string output_path;
std::vector<std::string> op_names;
std::vector<std::string> api_dirs;
std::string tf_prefix_dir;
std::string tf_root_dir;
std::string tf_output_dir;
explicit PathConfig() = default;
explicit PathConfig(const std::string& output_dir,
const std::string& source_dir,
const std::string& api_dir_list,
const std::vector<std::string> op_names);
};
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_PATH_CONFIG_H_
tensorflow/c/experimental/ops/gen/common/source_code.cc
+68
View File
@@ -0,0 +1,68 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/source_code.h"
#include <string>
#include "absl/log/log.h"
#include "absl/strings/ascii.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/strings/strip.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/stringpiece.h"
namespace tensorflow {
namespace generator {
std::string SourceCode::Render() const {
std::string code;
for (const Line& line : lines_) {
absl::StrAppend(&code, std::string(line.indent * spaces_per_indent_, ' '),
line.text, "\n");
}
return code;
}
void SourceCode::AddLineWithIndent(const std::string& line) {
ValidateAndAddLine(current_indent_, line);
}
void SourceCode::AddLineWithoutIndent(const std::string& line) {
ValidateAndAddLine(0, line);
}
void SourceCode::AddBlankLine() { ValidateAndAddLine(0, ""); }
void SourceCode::IncreaseIndent() { current_indent_++; }
void SourceCode::DecreaseIndent() { current_indent_--; }
void SourceCode::ValidateAndAddLine(int indent, const std::string& raw_line) {
absl::string_view line(raw_line);
bool had_trailing_newline = absl::ConsumeSuffix(&line, "\n");
if (absl::StrContains(line, '\n')) {
LOG(ERROR) << "Attempt to add multiple lines; bad formatting is likely.";
} else if (had_trailing_newline) {
LOG(WARNING) << "Superfluous trailing newline in '" << line << "'";
}
lines_.push_back(
{indent, std::string(absl::StripTrailingAsciiWhitespace(line))});
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/source_code.h
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_COMMON_SOURCE_CODE_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_SOURCE_CODE_H_
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
class SourceCode {
public:
std::string Render() const;
void SetSpacesPerIndent(int spaces_per_indent) {
spaces_per_indent_ = spaces_per_indent;
}
void AddLineWithIndent(const std::string& line);
void AddLineWithoutIndent(const std::string& line);
void AddBlankLine();
void IncreaseIndent();
void DecreaseIndent();
private:
struct Line {
int indent;
std::string text;
};
void ValidateAndAddLine(int indent_level, const std::string& raw_line);
int spaces_per_indent_ = 2;
int current_indent_ = 0;
std::vector<Line> lines_;
};
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_SOURCE_CODE_H_
tensorflow/c/experimental/ops/gen/common/view_util.cc
+43
View File
@@ -0,0 +1,43 @@
/* Copyright 2021 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/c/experimental/ops/gen/common/view_util.h"
#include <string>
#include <vector>
#include "absl/strings/str_join.h"
#include "absl/strings/substitute.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
std::string Call(const std::string& object, const std::string& method,
std::vector<std::string> arguments, const char* oper) {
return absl::Substitute("$0$1$2($3)", object, oper, method,
absl::StrJoin(arguments, ", "));
}
std::string Call(const std::string& function,
std::vector<std::string> arguments) {
return absl::Substitute("$0($1)", function, absl::StrJoin(arguments, ", "));
}
std::string Quoted(const std::string& s) {
return absl::Substitute("\"$0\"", s);
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/common/view_util.h
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_COMMON_VIEW_UTIL_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_VIEW_UTIL_H_
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
std::string Call(const std::string& function,
std::vector<std::string> arguments);
std::string Call(const std::string& object, const std::string& method,
std::vector<std::string> arguments, const char* oper = "->");
std::string Quoted(const std::string& s);
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_COMMON_VIEW_UTIL_H_
tensorflow/c/experimental/ops/gen/cpp/BUILD
+53
View File
@@ -0,0 +1,53 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_test",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
cc_library(
name = "cpp",
srcs = glob(
["*.cc"],
exclude = ["*_test.cc"],
),
hdrs = glob(["*.h"]),
visibility = ["//tensorflow/c/experimental/ops/gen:__pkg__"],
deps = [
"//tensorflow/c/experimental/ops/gen/common",
"//tensorflow/c/experimental/ops/gen/cpp/renderers",
"//tensorflow/c/experimental/ops/gen/cpp/views",
"//tensorflow/c/experimental/ops/gen/model",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"//tensorflow/core:op_gen_lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
tf_cc_test(
name = "cpp_generator_test",
size = "small",
srcs = ["cpp_generator_test.cc"],
data = ["//tensorflow/c/experimental/ops/gen/cpp/golden"],
deps = [
":cpp",
"//tensorflow/c/experimental/ops/gen/common",
"//tensorflow/c/experimental/ops/gen/cpp/renderers",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"@xla//xla/tsl/platform:status",
],
)
tensorflow/c/experimental/ops/gen/cpp/cpp_generator.cc
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/cpp_generator.h"
#include <string>
#include <vector>
#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
#include "tensorflow/c/experimental/ops/gen/common/source_code.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_config.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_file_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
#include "tensorflow/c/experimental/ops/gen/model/op_spec.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
CppGenerator::CppGenerator(cpp::CppConfig cpp_config, PathConfig path_config)
: controller_(path_config),
cpp_config_(cpp_config),
path_config_(path_config) {}
SourceCode CppGenerator::GenerateOneFile(
cpp::RendererContext::Mode mode) const {
SourceCode generated_code;
const std::vector<OpSpec> ops(controller_.GetModelOps());
std::vector<cpp::OpView> views(ops.begin(), ops.end());
cpp::RendererContext context{mode, generated_code, cpp_config_, path_config_};
cpp::CppFileRenderer(context, views).Render();
return generated_code;
}
SourceCode CppGenerator::HeaderFileContents() const {
return GenerateOneFile(cpp::RendererContext::kHeader);
}
SourceCode CppGenerator::SourceFileContents() const {
return GenerateOneFile(cpp::RendererContext::kSource);
}
std::string CppGenerator::HeaderFileName() const {
return io::JoinPath(path_config_.output_path, cpp_config_.unit + "_ops.h");
}
std::string CppGenerator::SourceFileName() const {
return io::JoinPath(path_config_.output_path, cpp_config_.unit + "_ops.cc");
}
void CppGenerator::WriteHeaderFile() const {
controller_.WriteFile(HeaderFileName(), HeaderFileContents());
}
void CppGenerator::WriteSourceFile() const {
controller_.WriteFile(SourceFileName(), SourceFileContents());
}
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/cpp_generator.h
+49
View File
@@ -0,0 +1,49 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_CPP_GENERATOR_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_CPP_GENERATOR_H_
#include "tensorflow/c/experimental/ops/gen/common/controller.h"
#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
#include "tensorflow/c/experimental/ops/gen/common/source_code.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_config.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
class CppGenerator {
public:
explicit CppGenerator(cpp::CppConfig cpp_config, PathConfig path_config);
SourceCode HeaderFileContents() const;
SourceCode SourceFileContents() const;
std::string HeaderFileName() const;
std::string SourceFileName() const;
void WriteHeaderFile() const;
void WriteSourceFile() const;
private:
SourceCode GenerateOneFile(cpp::RendererContext::Mode mode) const;
Controller controller_;
cpp::CppConfig cpp_config_;
PathConfig path_config_;
};
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_CPP_GENERATOR_H_
tensorflow/c/experimental/ops/gen/cpp/cpp_generator_test.cc
+83
View File
@@ -0,0 +1,83 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/cpp_generator.h"
#include <algorithm>
#include <string>
#include <vector>
#include "tensorflow/c/experimental/ops/gen/common/path_config.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_config.h"
#include "xla/tsl/platform/status.h"
#include "tensorflow/core/platform/env.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace {
TEST(CppGeneratorTest, typical_usage) {
std::string category = "testing";
std::string name_space = "tensorflow::ops";
std::string output_dir = "tensorflow/c/experimental/ops/gen/cpp/golden";
std::string source_dir = "tensorflow";
std::string api_dirs = "";
std::vector<std::string> ops = {
"Neg", // Simple unary Op
"MatMul", // 2 inputs & attrs with default values
"IdentityN", // Variadic input+output
"SparseSoftmaxCrossEntropyWithLogits", // 2 outputs
"AccumulatorApplyGradient", // 0 outputs
"VarHandleOp", // type, shape, list(string) attrs
"RestoreV2", // Variadic output-only, list(type) attr
};
cpp::CppConfig cpp_config(category, name_space);
PathConfig controller_config(output_dir, source_dir, api_dirs, ops);
CppGenerator generator(cpp_config, controller_config);
Env *env = Env::Default();
std::string golden_dir = io::JoinPath(testing::TensorFlowSrcRoot(),
controller_config.tf_output_dir);
std::string generated_header = generator.HeaderFileContents().Render();
std::string generated_source = generator.SourceFileContents().Render();
std::string expected_header;
std::string header_file_name =
io::JoinPath(golden_dir, "testing_ops.h.golden");
TF_CHECK_OK(ReadFileToString(env, header_file_name, &expected_header));
std::string expected_source;
std::string source_file_name =
io::JoinPath(golden_dir, "testing_ops.cc.golden");
TF_CHECK_OK(ReadFileToString(env, source_file_name, &expected_source));
// Remove carriage returns (for Windows)
expected_header.erase(
std::remove(expected_header.begin(), expected_header.end(), '\r'),
expected_header.end());
expected_source.erase(
std::remove(expected_source.begin(), expected_source.end(), '\r'),
expected_source.end());
EXPECT_EQ(expected_header, generated_header);
EXPECT_EQ(expected_source, generated_source);
}
} // namespace
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/golden/BUILD
+10
View File
@@ -0,0 +1,10 @@
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:public"],
licenses = ["notice"],
)
filegroup(
name = "golden",
data = glob(["*.golden"]),
)
tensorflow/c/experimental/ops/gen/cpp/golden/testing_ops.cc.golden
+150
View File
@@ -0,0 +1,150 @@
/* Copyright 2021 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.
==============================================================================*/
// This file is MACHINE GENERATED! Do not edit.
#include "tensorflow/c/experimental/ops/gen/cpp/golden/testing_ops.h"
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/core/framework/types.h" // NOLINT
#include <cstring> // NOLINT
#include "tensorflow/c/eager/abstract_operation.h"
#include "tensorflow/c/eager/tracing_utils.h"
#include "tensorflow/core/platform/errors.h" // NOLINT
using tensorflow::tracing::MaybeSetOpName;
namespace tensorflow {
namespace ops {
// Op: Neg()
// Summary:
//
// Description:
absl::Status Neg(AbstractContext* ctx, AbstractTensorHandle* const x, AbstractTensorHandle** y, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("Neg", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(x));
int num_retvals = 1;
return op_ptr->Execute(absl::MakeSpan(y, 1), &num_retvals);
}
// Op: MatMul()
// Summary:
//
// Description:
absl::Status MatMul(AbstractContext* ctx, AbstractTensorHandle* const a, AbstractTensorHandle* const b, AbstractTensorHandle** product, bool transpose_a, bool transpose_b, bool grad_a, bool grad_b, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("MatMul", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(a));
TF_RETURN_IF_ERROR(op_ptr->AddInput(b));
TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_a", transpose_a));
TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("transpose_b", transpose_b));
TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_a", grad_a));
TF_RETURN_IF_ERROR(op_ptr->SetAttrBool("grad_b", grad_b));
int num_retvals = 1;
return op_ptr->Execute(absl::MakeSpan(product, 1), &num_retvals);
}
// Op: IdentityN()
// Summary:
//
// Description:
absl::Status IdentityN(AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> input, absl::Span<AbstractTensorHandle*> output, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("IdentityN", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInputList(input));
int num_retvals = output.size();
return op_ptr->Execute(output, &num_retvals);
}
// Op: SparseSoftmaxCrossEntropyWithLogits()
// Summary:
//
// Description:
absl::Status SparseSoftmaxCrossEntropyWithLogits(AbstractContext* ctx, AbstractTensorHandle* const features, AbstractTensorHandle* const labels, AbstractTensorHandle** loss, AbstractTensorHandle** backprop, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("SparseSoftmaxCrossEntropyWithLogits", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(features));
TF_RETURN_IF_ERROR(op_ptr->AddInput(labels));
int num_retvals = 2;
AbstractTensorHandle* temp_outputs[2] = {nullptr};
absl::Status status = op_ptr->Execute(temp_outputs, &num_retvals);
if (status.ok()) {
*loss = temp_outputs[0];
*backprop = temp_outputs[1];
}
return status;
}
// Op: AccumulatorApplyGradient()
// Summary:
//
// Description:
absl::Status AccumulatorApplyGradient(AbstractContext* ctx, AbstractTensorHandle** handle, AbstractTensorHandle* const local_step, AbstractTensorHandle* const gradient, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("AccumulatorApplyGradient", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(handle));
TF_RETURN_IF_ERROR(op_ptr->AddInput(local_step));
TF_RETURN_IF_ERROR(op_ptr->AddInput(gradient));
int num_retvals = 0;
std::vector<AbstractTensorHandle*> dummy_outputs;
return op_ptr->Execute(absl::MakeSpan(dummy_outputs), &num_retvals);
}
// Op: VarHandleOp()
// Summary:
//
// Description:
absl::Status VarHandleOp(AbstractContext* ctx, AbstractTensorHandle** resource, DataType dtype, const PartialTensorShape shape, const char* container, const char* shared_name, const char* debug_name, absl::Span<std::string const> allowed_devices, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("VarHandleOp", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->SetAttrString("container", container, strlen(container)));
TF_RETURN_IF_ERROR(op_ptr->SetAttrString("shared_name", shared_name, strlen(shared_name)));
TF_RETURN_IF_ERROR(op_ptr->SetAttrString("debug_name", debug_name, strlen(debug_name)));
TF_RETURN_IF_ERROR(op_ptr->SetAttrType("dtype", dtype));
TF_RETURN_IF_ERROR(op_ptr->SetAttrShape("shape", shape));
TF_RETURN_IF_ERROR(op_ptr->SetAttrStringList("allowed_devices", allowed_devices));
int num_retvals = 1;
return op_ptr->Execute(absl::MakeSpan(resource, 1), &num_retvals);
}
// Op: RestoreV2()
// Summary:
//
// Description:
absl::Status RestoreV2(AbstractContext* ctx, AbstractTensorHandle* const prefix, AbstractTensorHandle* const tensor_names, AbstractTensorHandle* const shape_and_slices, absl::Span<AbstractTensorHandle*> tensors, absl::Span<DataType> dtypes, const char* name, const char* raw_device_name) {
AbstractOperationPtr op_ptr(ctx->CreateOperation());
TF_RETURN_IF_ERROR(op_ptr->Reset("RestoreV2", raw_device_name));
TF_RETURN_IF_ERROR(MaybeSetOpName(op_ptr.get(), name));
TF_RETURN_IF_ERROR(op_ptr->AddInput(prefix));
TF_RETURN_IF_ERROR(op_ptr->AddInput(tensor_names));
TF_RETURN_IF_ERROR(op_ptr->AddInput(shape_and_slices));
TF_RETURN_IF_ERROR(op_ptr->SetAttrTypeList("dtypes", dtypes.data(), dtypes.length()));
int num_retvals = tensors.size();
return op_ptr->Execute(tensors, &num_retvals);
}
} // namespace ops
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/golden/testing_ops.h.golden
+54
View File
@@ -0,0 +1,54 @@
/* Copyright 2021 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.
==============================================================================*/
// This file is MACHINE GENERATED! Do not edit.
#ifndef TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_GOLDEN_TESTING_OPS_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_GOLDEN_TESTING_OPS_H_
#include "absl/status/status.h"
#include "absl/types/span.h"
#include "tensorflow/c/eager/abstract_context.h"
#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/core/framework/types.h" // NOLINT
namespace tensorflow {
namespace ops {
//
absl::Status Neg(AbstractContext* ctx, AbstractTensorHandle* const x, AbstractTensorHandle** y, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status MatMul(AbstractContext* ctx, AbstractTensorHandle* const a, AbstractTensorHandle* const b, AbstractTensorHandle** product, bool transpose_a = false, bool transpose_b = false, bool grad_a = false, bool grad_b = false, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status IdentityN(AbstractContext* ctx, absl::Span<AbstractTensorHandle* const> input, absl::Span<AbstractTensorHandle*> output, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status SparseSoftmaxCrossEntropyWithLogits(AbstractContext* ctx, AbstractTensorHandle* const features, AbstractTensorHandle* const labels, AbstractTensorHandle** loss, AbstractTensorHandle** backprop, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status AccumulatorApplyGradient(AbstractContext* ctx, AbstractTensorHandle** handle, AbstractTensorHandle* const local_step, AbstractTensorHandle* const gradient, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status VarHandleOp(AbstractContext* ctx, AbstractTensorHandle** resource, DataType dtype, const PartialTensorShape shape, const char* container = "", const char* shared_name = "", const char* debug_name = "", absl::Span<std::string const> allowed_devices = {}, const char* name = nullptr, const char* raw_device_name = nullptr);
//
absl::Status RestoreV2(AbstractContext* ctx, AbstractTensorHandle* const prefix, AbstractTensorHandle* const tensor_names, AbstractTensorHandle* const shape_and_slices, absl::Span<AbstractTensorHandle*> tensors, absl::Span<DataType> dtypes, const char* name = nullptr, const char* raw_device_name = nullptr);
} // namespace ops
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_GOLDEN_TESTING_OPS_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/BUILD
+52
View File
@@ -0,0 +1,52 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load(
"//tensorflow:tensorflow.bzl",
"tf_cc_tests",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = ["//visibility:private"],
licenses = ["notice"],
)
cc_library(
name = "renderers",
srcs = glob(
["*.cc"],
exclude = ["*_test.cc"],
),
hdrs = glob(["*.h"]),
visibility = [
"//tensorflow/c/experimental/ops/gen:__pkg__",
"//tensorflow/c/experimental/ops/gen/cpp:__pkg__",
],
deps = [
"//tensorflow/c/experimental/ops/gen/common",
"//tensorflow/c/experimental/ops/gen/cpp/views",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:lib_internal",
"@com_google_absl//absl/log",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
tf_cc_tests(
name = "renderer_tests",
size = "small",
srcs = glob(
["*_test.cc"],
),
deps = [
":renderers",
"//tensorflow/c/experimental/ops/gen/common",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
"//tensorflow/core:test",
"//tensorflow/core:test_main",
"//tensorflow/core:testlib",
"//tensorflow/core/platform:types",
],
)
tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_config.cc
+34
View File
@@ -0,0 +1,34 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/cpp_config.h"
#include <string>
#include "absl/strings/ascii.h"
#include "absl/strings/str_split.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
CppConfig::CppConfig(const std::string& category, const std::string& name_space)
: category(category),
unit(absl::AsciiStrToLower(category)),
namespaces(absl::StrSplit(name_space, "::")) {}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_config.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_CONFIG_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_CONFIG_H_
#include <vector>
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
struct CppConfig {
std::string category;
std::string unit;
std::vector<std::string> namespaces;
explicit CppConfig() = default;
explicit CppConfig(const std::string& category,
const std::string& name_space = "tensorflow::ops");
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_CONFIG_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_file_renderer.cc
+85
View File
@@ -0,0 +1,85 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/cpp_file_renderer.h"
#include <vector>
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/op_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
namespace tensorflow {
namespace generator {
namespace cpp {
static const char *copyright =
R"(
/* Copyright 2021 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.
==============================================================================*/
)";
static const char *machine_generated =
"// This file is MACHINE GENERATED! Do not edit.";
CppFileRenderer::CppFileRenderer(RendererContext context,
const std::vector<OpView> &ops)
: Renderer(context),
guard_(context),
name_space_(context),
includes_(context),
ops_(ops) {}
void CppFileRenderer::Render() {
CodeLines(copyright);
BlankLine();
CodeLine(machine_generated);
BlankLine();
if (context_.mode == RendererContext::kHeader) {
guard_.Open();
} else {
includes_.SelfHeader();
}
includes_.Headers(ops_);
name_space_.Open();
BlankLine();
for (const OpView &op : ops_) {
OpRenderer(context_, op).Render();
}
name_space_.Close();
if (context_.mode == RendererContext::kHeader) {
guard_.Close();
}
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/cpp_file_renderer.h
+48
View File
@@ -0,0 +1,48 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_FILE_RENDERER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_FILE_RENDERER_H_
#include <vector>
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/guard_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/include_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/namespace_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
namespace tensorflow {
namespace generator {
namespace cpp {
class CppFileRenderer : public Renderer {
public:
explicit CppFileRenderer(RendererContext context,
const std::vector<OpView> &ops);
void Render();
private:
GuardRenderer guard_;
NamespaceRenderer name_space_;
IncludeRenderer includes_;
std::vector<OpView> ops_;
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_CPP_FILE_RENDERER_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/guard_renderer.cc
+53
View File
@@ -0,0 +1,53 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/guard_renderer.h"
#include <algorithm>
#include <string>
#include "tensorflow/c/experimental/ops/gen/common/case_format.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
GuardRenderer::GuardRenderer(RendererContext context) : Renderer(context) {
std::string self_path = io::JoinPath(context_.path_config.tf_root_dir,
context_.path_config.tf_output_dir,
context_.cpp_config.unit + "_ops.h");
std::string with_underscores(self_path);
std::replace(with_underscores.begin(), with_underscores.end(), '/', '_');
std::replace(with_underscores.begin(), with_underscores.end(), '.', '_');
guard_ = toUpperSnake(with_underscores) + "_";
}
void GuardRenderer::Open() {
CodeLine("#ifndef $0", guard_);
CodeLine("#define $0", guard_);
BlankLine();
}
void GuardRenderer::Close() {
BlankLine();
CodeLine("#endif // $0", guard_);
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/guard_renderer.h
+41
View File
@@ -0,0 +1,41 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_GUARD_RENDERER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_GUARD_RENDERER_H_
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
class GuardRenderer : public Renderer {
public:
explicit GuardRenderer(RendererContext context);
void Open();
void Close();
private:
std::string guard_;
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_GUARD_RENDERER_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/include_renderer.cc
+98
View File
@@ -0,0 +1,98 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/include_renderer.h"
#include <string>
#include <vector>
#include "absl/strings/match.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_argument_view.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
#include "tensorflow/core/platform/path.h"
namespace tensorflow {
namespace generator {
namespace cpp {
IncludeRenderer::IncludeRenderer(RendererContext context) : Renderer(context) {}
void IncludeRenderer::SelfHeader() {
Include(SelfHeaderPath());
BlankLine();
}
std::string IncludeRenderer::SelfHeaderPath() const {
return io::JoinPath(context_.path_config.tf_root_dir,
context_.path_config.tf_output_dir,
context_.cpp_config.unit + "_ops.h");
}
void IncludeRenderer::Include(const std::string& tf_file_path) {
CodeLine("#include \"$0\"",
io::JoinPath(context_.path_config.tf_prefix_dir, tf_file_path));
}
void IncludeRenderer::Headers(const std::vector<OpView>& ops) {
Include(
"absl"
"/status/status.h");
bool needs_span = false;
if (context_.mode == RendererContext::kSource) {
needs_span = true;
} else {
for (const OpView& op : ops) {
for (const OpArgumentView& arg : op.AllArguments()) {
if (absl::StrContains(arg.Declaration(), "absl::Span")) {
needs_span = true;
break;
}
}
if (needs_span) break;
}
}
if (needs_span) {
Include(
"absl"
"/types/span.h");
}
Include("tensorflow/c/eager/abstract_context.h");
Include("tensorflow/c/eager/abstract_tensor_handle.h");
if (context_.cpp_config.unit == "resource_variable") {
Include("tensorflow/core/framework/tensor_shape.h");
}
CodeLine("#include \"$0\" // NOLINT",
io::JoinPath(context_.path_config.tf_prefix_dir,
"tensorflow/core/framework/types.h"));
if (context_.cpp_config.unit == "resource_variable") {
CodeLine("#include <string>");
}
if (context_.mode == RendererContext::kSource) {
CodeLine("#include <cstring> // NOLINT");
Include("tensorflow/c/eager/abstract_operation.h");
Include("tensorflow/c/eager/tracing_utils.h");
CodeLine("#include \"$0\" // NOLINT",
io::JoinPath(context_.path_config.tf_prefix_dir,
"tensorflow/core/platform/errors.h"));
BlankLine();
Statement("using tensorflow::tracing::MaybeSetOpName");
}
BlankLine();
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/include_renderer.h
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_INCLUDE_RENDERER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_INCLUDE_RENDERER_H_
#include <vector>
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
class IncludeRenderer : public Renderer {
public:
explicit IncludeRenderer(RendererContext context);
std::string SelfHeaderPath() const;
void SelfHeader();
void Headers(const std::vector<OpView>& ops);
private:
void Include(const std::string& tf_file_path);
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_INCLUDE_RENDERER_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/namespace_renderer.cc
+45
View File
@@ -0,0 +1,45 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/namespace_renderer.h"
#include <string>
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
NamespaceRenderer::NamespaceRenderer(RendererContext context)
: Renderer(context) {}
void NamespaceRenderer::Open() {
for (const std::string& ns : context_.cpp_config.namespaces) {
CodeLine("namespace " + ns + " {");
}
}
void NamespaceRenderer::Close() {
for (auto it = context_.cpp_config.namespaces.rbegin();
it != context_.cpp_config.namespaces.rend(); ++it) {
CodeLine("} // namespace " + *it);
}
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/namespace_renderer.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_NAMESPACE_RENDERER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_NAMESPACE_RENDERER_H_
#include <vector>
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/core/platform/types.h"
namespace tensorflow {
namespace generator {
namespace cpp {
class NamespaceRenderer : public Renderer {
public:
explicit NamespaceRenderer(RendererContext context);
void Open();
void Close();
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_NAMESPACE_RENDERER_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/op_comment_renderer.cc
+46
View File
@@ -0,0 +1,46 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/op_comment_renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
namespace tensorflow {
namespace generator {
namespace cpp {
OpCommentRenderer::OpCommentRenderer(RendererContext context, OpView op)
: Renderer(context), op_(op) {}
void OpCommentRenderer::Render() {
if (context_.mode == RendererContext::kHeader) {
// Add a short 1-line comment to the header files.
CommentLine(op_.Summary());
return;
}
CommentLine("Op: $0()", op_.FunctionName());
CommentLine("Summary: $0", op_.Summary());
CommentLine("");
CommentLine("Description:");
for (const auto& line : op_.Description()) {
CommentLine(" $0", line);
}
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow
tensorflow/c/experimental/ops/gen/cpp/renderers/op_comment_renderer.h
+40
View File
@@ -0,0 +1,40 @@
/* Copyright 2021 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_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_OP_COMMENT_RENDERER_H_
#define TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_OP_COMMENT_RENDERER_H_
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
namespace tensorflow {
namespace generator {
namespace cpp {
class OpCommentRenderer : public Renderer {
public:
explicit OpCommentRenderer(RendererContext context, OpView op);
void Render();
private:
OpView op_;
};
} // namespace cpp
} // namespace generator
} // namespace tensorflow
#endif // TENSORFLOW_C_EXPERIMENTAL_OPS_GEN_CPP_RENDERERS_OP_COMMENT_RENDERER_H_
tensorflow/c/experimental/ops/gen/cpp/renderers/op_implementation_renderer.cc
+102
View File
@@ -0,0 +1,102 @@
/* Copyright 2021 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/c/experimental/ops/gen/cpp/renderers/op_implementation_renderer.h"
#include "tensorflow/c/experimental/ops/gen/common/view_util.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer.h"
#include "tensorflow/c/experimental/ops/gen/cpp/renderers/renderer_context.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/arg_view.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/attr_view.h"
#include "tensorflow/c/experimental/ops/gen/cpp/views/op_view.h"
namespace tensorflow {
namespace generator {
namespace cpp {
OpImplementationRenderer::OpImplementationRenderer(RendererContext context,
OpView op)
: Renderer(context), op_(op) {}
void OpImplementationRenderer::Render() {
RenderInitialization();
if (op_.IsListOp()) {
RenderExecutionListOp();
} else if (op_.NumOutputs() == 0) {
RenderExecutionZeroOutputs();
} else if (op_.NumOutputs() == 1) {
RenderExecutionSingleOutput();
} else {
RenderExecutionMultipleOutputs();
}
}
void OpImplementationRenderer::RenderInitialization() {
// Create Op variable and initialize it
Statement("AbstractOperationPtr $0(ctx->CreateOperation())",
op_.VariableName());
TFStatement(Call(op_.VariableName(), "Reset",
{op_.OpNameString(), "raw_device_name"}));
TFStatement(Call("MaybeSetOpName", {op_.VariableName() + ".get()", "name"}));
// Set each input
for (const ArgView& ar : op_.Inputs()) {
TFStatement(Call(op_.VariableName(), ar.SetterMethod(), ar.SetterArgs()));
}
// Set each attribute
for (const AttrView& ar : op_.Attributes()) {
TFStatement(Call(op_.VariableName(), ar.SetterMethod(), ar.SetterArgs()));
}
}
void OpImplementationRenderer::RenderExecutionListOp() {
ArgView output_arg = op_.OnlyOutput();
Statement("int num_retvals = $0.size()", output_arg.VariableName());
Statement("return " + Call(op_.VariableName(), "Execute",
{output_arg.VariableName(), "&num_retvals"}));
}
void OpImplementationRenderer::RenderExecutionSingleOutput() {
ArgView output_arg = op_.OnlyOutput();
Statement("int num_retvals = 1");
Statement("return $0->Execute(absl::MakeSpan($1, 1), &num_retvals)",
op_.VariableName(), output_arg.VariableName());
}
void OpImplementationRenderer::RenderExecutionMultipleOutputs() {
Statement("int num_retvals = $0", op_.NumOutputs());
Statement("AbstractTensorHandle* temp_outputs[$0] = {nullptr}",
op_.NumOutputs());
Statement("absl::Status status = $0->Execute(temp_outputs, &num_retvals)",
op_.VariableName());
BlockOpen("if (status.ok())");
for (const ArgView& arg : op_.Outputs()) {
Statement("*$0 = temp_outputs[$1]", arg.VariableName(), arg.Position());
}
BlockClose();
Statement("return status");
}
void OpImplementationRenderer::RenderExecutionZeroOutputs() {
Statement("int num_retvals = 0");
Statement("std::vector<AbstractTensorHandle*> dummy_outputs");
Statement("return $0->Execute(absl::MakeSpan(dummy_outputs), &num_retvals)",
op_.VariableName());
}
} // namespace cpp
} // namespace generator
} // namespace tensorflow

Some files were not shown because too many files have changed in this diff Show More