wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 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/compiler/mlir/lite/experimental/remat/BUILD
+44
View File
@@ -0,0 +1,44 @@
load("@rules_cc//cc:cc_library.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_cc_test")
load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
# copybara:uncomment package(default_applicable_licenses = ["//tensorflow:LICENSE"])
cc_library(
name = "rematerializer",
srcs = ["rematerializer.cc"],
hdrs = ["rematerializer.h"],
deps = [
],
)
tf_cc_test(
name = "rematerializer_test",
size = "small",
srcs = ["rematerializer_test.cc"],
deps = [
":rematerializer",
"@com_google_googletest//:gtest_main",
],
)
cc_library(
name = "metadata_util",
srcs = ["metadata_util.cc"],
hdrs = ["metadata_util.h"],
compatible_with = get_compatible_with_portable(),
visibility = ["//visibility:public"],
deps = [
"//tensorflow/compiler/mlir/lite:control_edges",
],
)
tf_cc_test(
name = "metadata_util_test",
size = "small",
srcs = ["metadata_util_test.cc"],
deps = [
":metadata_util",
"@com_google_googletest//:gtest_main",
],
)
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.cc
+124
View File
@@ -0,0 +1,124 @@
/* 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/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstddef>
#include <cstdint>
#include <string>
#include <utility>
#include <vector>
namespace {
// We serialize unsigneds as protobuf varints, i.e., in chunks of 7 bits each.
constexpr int kMod = (1 << 7);
void Serialize(std::string* out, uint32_t value) {
for (; value >= kMod; value /= kMod) {
out->push_back(value % kMod + kMod);
}
out->push_back(value);
}
bool Parse(const char** data, size_t* size, uint32_t* out) {
*out = 0;
uint32_t mul = 1;
for (bool done = false; !done;
mul *= kMod, done = !(**data & kMod), ++*data, --*size) {
if (*size == 0) {
return false;
}
*out += static_cast<unsigned char>(**data) % kMod * mul;
}
return true;
}
// Signed ints are zigzag-encoded as unsigned varints, [..., -2, -1, 0, 1, 2,
// ...] ->
// [..., 3, 1, 0, 2, 4, ...].
void Serialize(std::string* out, int32_t value) {
Serialize(out, static_cast<uint32_t>(
value < 0 ? static_cast<uint32_t>(-(value + 1)) * 2 + 1
: static_cast<uint32_t>(value) * 2));
}
bool Parse(const char** data, size_t* size, int32_t* out) {
uint32_t value = 0;
if (!Parse(data, size, &value)) {
return false;
}
const int32_t magnitude = value / 2;
*out = (value % 2) ? (-magnitude - 1) : magnitude;
return true;
}
// Pairs are serialized as the concatenation of their elements' serialization.
template <class First, class Second>
void Serialize(std::string* out, const std::pair<First, Second>& in) {
Serialize(out, in.first);
Serialize(out, in.second);
}
template <class First, class Second>
bool Parse(const char** data, size_t* size, std::pair<First, Second>* out) {
return Parse(data, size, &(out->first)) && Parse(data, size, &(out->second));
}
// Vectors are serialized as the concetation of the serialization of their size
// and the the serializations of their elements.
template <class Value>
void Serialize(std::string* out, const std::vector<Value>& in) {
Serialize(out, static_cast<uint32_t>(in.size()));
for (const auto& val : in) {
Serialize(out, val);
}
}
template <class T>
bool Parse(const char** data, size_t* size, std::vector<T>* out) {
uint32_t num_elems = 0;
if (!Parse(data, size, &num_elems)) {
return false;
}
out->assign(num_elems, T{});
for (auto& elem : *out) {
if (!Parse(data, size, &elem)) {
return false;
}
}
return true;
}
} // namespace
namespace tflite {
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in) {
std::string out;
Serialize(&out, kModelControlDependenciesMetadataVersion);
Serialize(&out, in);
return out;
}
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out) {
out->clear();
uint32_t version = 0;
return Parse(&data, &size, &version) &&
(version == kModelControlDependenciesMetadataVersion) &&
Parse(&data, &size, out) && (size == 0);
}
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util.h
+62
View File
@@ -0,0 +1,62 @@
/* 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.
==============================================================================*/
/// \file
///
/// Functions for serializiation/deserialization of control dependency
/// information to/from model metadata.
///
#ifndef TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
#include <cstddef>
#include <cstdint>
#include <string>
#include <vector>
#include "tensorflow/compiler/mlir/lite/utils/control_edges.h"
namespace tflite {
/// Control dependencies for the model is the collection of control dependencies
/// for its subgraphs.
using ModelControlDependencies = std::vector<ControlEdges>;
/// Serializes `in` into the returned string. The result is parseable with
/// ParseModelControlDependencies.
std::string SerializeModelControlDependencies(
const ModelControlDependencies& in);
/// Deserializes `*out` from a character buffer of size `size` at `data`.
/// Returns true iff successful. `*out` needn't be empty before invocation.
/// When returning false, `*out`'s state is undefined.
bool ParseModelControlDependencies(const char* data, size_t size,
ModelControlDependencies* out);
/// The key under which to store the serialized control dependencies in the
/// model's metadata.
constexpr char kModelControlDependenciesMetadataKey[] =
"model_control_dependencies";
/// To allow future changes to the format, serialized control dependency data
/// will contain a version; this constant is the version that will be used for
/// serialization. For deserialization, past versions should remain parseable.
constexpr uint32_t kModelControlDependenciesMetadataVersion = 1;
inline constexpr char kModelUseStablehloTensorKey[] = "keep_stablehlo_constant";
} // namespace tflite
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_METADATA_UTIL_H_
tensorflow/compiler/mlir/lite/experimental/remat/metadata_util_test.cc
+55
View File
@@ -0,0 +1,55 @@
/* 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/compiler/mlir/lite/experimental/remat/metadata_util.h"
#include <cstdint>
#include <limits>
#include <string>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace tflite {
namespace {
class MetadataSerializerTest : public ::testing::Test {
protected:
static constexpr auto kHuge = std::numeric_limits<int32_t>::max();
static constexpr auto kTiny = std::numeric_limits<int32_t>::min();
std::string RoundTrip(const ModelControlDependencies &in) const {
ModelControlDependencies out = {{{-1, -1}}};
const std::string serialized =
tflite::SerializeModelControlDependencies(in);
return tflite::ParseModelControlDependencies(serialized.data(),
serialized.size(), &out)
? (out == in) ? "ok" : "mismatch"
: "malformed";
}
};
TEST_F(MetadataSerializerTest, nothing) { EXPECT_THAT(RoundTrip({}), "ok"); }
TEST_F(MetadataSerializerTest, something) {
EXPECT_THAT(
RoundTrip({{{1, 2}, {2, 3}, {4, 5}},
{},
{{kHuge, kTiny}, {kTiny, kHuge}, {kHuge - 1, kTiny + 1}},
{{1, 0}}}),
"ok");
}
} // namespace
} // namespace tflite
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.cc
+381
View File
@@ -0,0 +1,381 @@
/* 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/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <map>
#include <tuple>
#include <utility>
#include <vector>
namespace mlir {
namespace TFL {
namespace {
// Helper functions for sorted + deduped int vectors.
std::tuple<std::vector<int>::iterator, bool> Find(const int item,
std::vector<int>& items) {
const auto iter = std::lower_bound(items.begin(), items.end(), item);
return std::make_tuple(iter, iter != items.end() && *iter == item);
}
void Insert(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (!found) items.insert(iter, item);
}
void Erase(const int item, std::vector<int>& items) {
const auto [iter, found] = Find(item, items);
if (found) items.erase(iter);
}
} // namespace
int Rematerializer::AddOperation(const bool is_stateful) {
operations_.emplace_back();
operations_.back().is_stateful = is_stateful;
return operations_.size() - 1;
}
int Rematerializer::AddTensor(const SizeT size) {
tensors_.emplace_back();
tensors_.back().size = size;
return tensors_.size() - 1;
}
void Rematerializer::DelUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
// Was the dependence to be deleted the first/last (or both) use of this
// tensor?
const bool was_first_use =
(!tensor.operations.empty() && ioperation == tensor.first_use());
const bool was_last_use =
(!tensor.operations.empty() && ioperation == tensor.last_use());
Erase(ioperation, tensor.operations);
Erase(itensor, operation.tensors);
if (was_first_use) {
operation.alloc -= size;
if (!was_last_use) {
operations_[tensor.first_use()].alloc += size;
}
}
if (was_last_use) {
operation.dealloc -= size;
if (!was_first_use) {
operations_[tensor.last_use()].dealloc += size;
}
}
}
void Rematerializer::AddUse(const int ioperation, const int itensor) {
auto& tensor = tensors_[itensor];
auto& operation = operations_[ioperation];
const auto& size = tensor.size;
const bool will_be_first_use =
tensor.operations.empty() || ioperation < tensor.first_use();
const bool will_be_last_use =
tensor.operations.empty() || ioperation > tensor.last_use();
if (will_be_first_use) {
operation.alloc += size;
if (!will_be_last_use) {
operations_[tensor.first_use()].alloc -= size;
}
}
if (will_be_last_use) {
operation.dealloc += size;
if (!will_be_first_use) {
operations_[tensor.last_use()].dealloc -= size;
}
}
Insert(ioperation, tensor.operations);
Insert(itensor, operation.tensors);
}
Rematerializer::SizeT Rematerializer::MaxSavings(const int begin, const int end,
const int peak_loc) const {
SizeT max_savings = 0;
// We're looking at the outputs of all operators in [`begin`, `end`). If an
// output is alive after `peak_loc`, rematerializing the operator anywhere
// after `peak_loc` would lower the memory profile at `peak_loc`.
for (int ioperation = begin; ioperation != end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const Tensor& tensor = tensors_[itensor];
tensor.first_use() == ioperation /* output */ &&
tensor.last_use() > peak_loc /* used later */) {
max_savings += tensor.size;
}
}
}
return max_savings;
}
std::tuple<Rematerializer::SizeT, Rematerializer::RematSpec>
Rematerializer::FindBestRemat(const SizeT min_savings, const int begin_len,
const int end_len) const {
const auto peak = GetPeakMemory();
SizeT best_peak_mem = peak.size;
RematSpec best_remat = {};
for (int len = begin_len; len < end_len; ++len) {
std::vector<std::tuple<SizeT, int, int>> pre_screen;
for (int begin = 0, end = begin + len; end <= peak.op_index;
++begin, ++end) {
if (!std::any_of(operations_.begin() + begin, operations_.begin() + end,
[](const Operation& s) { return s.is_stateful; })) {
if (const auto max_savings = MaxSavings(begin, end, peak.op_index);
max_savings >= min_savings) {
pre_screen.emplace_back(max_savings, begin, end);
}
}
}
std::sort(pre_screen.begin(), pre_screen.end());
for (; !pre_screen.empty(); pre_screen.pop_back()) {
const auto& [max_savings, begin, end] = pre_screen.back();
const auto insert_before = FindBestRematPoint(begin, end, peak.op_index);
if (insert_before == operations_.size()) {
continue;
}
const RematSpec this_remat = {begin, end, insert_before};
if (const auto new_peak = GetPeakMemory(this_remat);
new_peak.size < best_peak_mem &&
peak.size >= new_peak.size + min_savings) {
best_peak_mem = new_peak.size;
best_remat = this_remat;
}
// If the actual savings achieved is bigger than the maximal savings that
// can be possibly achieved, leave early.
if (peak.size >= max_savings + best_peak_mem) {
break;
}
}
// We already found one savings for this length.
if (peak.size >= min_savings + best_peak_mem) {
break;
}
}
return std::make_tuple(best_peak_mem, best_remat);
}
std::vector<Rematerializer::MemSpec> Rematerializer::GetDeltas(
const RematSpec& remat) const {
// We're cloning operations [remat.begin, remat.end) at position
// remat.insert. We store changes to the alloc/dealloc sizes due to the
// insertion in a vector `delta`: A change `c_alloc` of `operations_[i].alloc`
// as `delta[i] += c_alloc`, and a change `c_dealloc` of
// `operations_[i].dealloc` as `delta[i+1] -= c_dealloc`.
std::vector<MemSpec> deltas;
if (remat.begin == remat.end) {
return deltas;
}
// Helper lambda: converts an operation index in the original range to its
// equivalent in the cloned range.
const auto source_to_target = [&](int i) {
return i + (remat.insert - remat.begin);
};
// Helper struct: bundles first and last use of a tensor within
// a contiguous range of operations.
struct TensorUse {
int first_use;
int last_use;
};
// For all tensors in the operation range to be cloned, store their first and
// last use within that range. This will be needed below to decide whether a
// tensor's life will be extended, or if it is to be rematerialized.
std::map<int, TensorUse> source_uses;
for (int ioperation = remat.begin; ioperation < remat.end; ++ioperation) {
const auto& operation = operations_[ioperation];
for (const int itensor : operation.tensors) {
// insertion will only happen for the first use.
const auto [iter, inserted] = source_uses.emplace(
itensor,
TensorUse{/*first_use=*/ioperation, /*last_use=*/ioperation});
if (!inserted) {
// Otherwise, update last_use.
iter->second.last_use = ioperation;
}
}
}
deltas.reserve(2 * source_uses.size());
for (const auto& [itensor, source] : source_uses) {
auto& tensor = tensors_[itensor];
const TensorUse global = {tensor.first_use(), tensor.last_use()};
auto add_alloc = [&](int pos) { deltas.emplace_back(pos, tensor.size); };
auto add_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, -tensor.size);
};
auto del_dealloc = [&](int pos) {
deltas.emplace_back(pos + 1, tensor.size);
};
if (global.first_use < remat.begin) {
// The tensor is created before the source range.
if (global.last_use < remat.insert) {
// It currently gets deallocated before the newly inserted range, so we
// need to extend its lifetime: It will now be deallocated at its last
// use in the inserted range.
del_dealloc(global.last_use);
add_dealloc(source_to_target(source.last_use));
}
} else {
// Tensor is created in the source range. It will be rematerialized in the
// cloned range.
add_alloc(source_to_target(source.first_use));
if (global.last_use < remat.insert) {
// The last use of the original tensor is before the target range, so
// the lifetime of the rematerialized tensor ends with the target range.
add_dealloc(source_to_target(source.last_use));
} else {
// There are uses of the original tensor after the cloned range. They
// can be replaced with uses of the rematerialized tensor, and the
// original tensor can be deallocated after its last use before the
// rematerialization.
add_dealloc(*std::partition_point(
tensor.operations.rbegin(), tensor.operations.rend(),
[&](int i) { return i >= remat.insert; }));
}
}
}
std::sort(deltas.begin(), deltas.end(), ByOpIndex);
return deltas;
}
Rematerializer::MemProfile Rematerializer::GetMemProfile(
const RematSpec& remat) const {
const auto num_inserted = remat.end - remat.begin;
std::vector<SizeT> profile(operations_.size() + num_inserted);
MapMem([&](const MemSpec& m) { profile[m.op_index] = m.size; }, remat);
return profile;
}
Rematerializer::MemSpec Rematerializer::GetPeakMemory(
const RematSpec& remat) const {
MemSpec peak;
MapMem([&](const MemSpec& m) { peak = std::max(m, peak, BySize); }, remat);
return peak;
}
int Rematerializer::FindBestRematPoint(const int begin, const int end,
const int peak_loc) const {
// All tensors necessarily have their last use before or at the last operator
// (which is the yield operation of a function), so an unchanged best ==
// operations_.size() value means that there is no valid configuration.
int best = operations_.size();
// All tensors whose first use is in [begin, end) and that are not destroyed
// until after peak_loc.
for (int ioperation = begin; ioperation < end; ++ioperation) {
for (const int itensor : operations_[ioperation].tensors) {
if (const auto& tensor = tensors_[itensor];
tensor.first_use() >= begin && tensor.first_use() < end &&
tensor.last_use() > peak_loc) {
for (const int ioperation : tensor.operations) {
// We return the earliest dependence on any output tensor.
if (ioperation > peak_loc && ioperation < best) {
best = ioperation;
break;
}
}
}
}
}
return best;
}
void Rematerializer::Remat(const RematSpec& remat) {
const int num_inserted = remat.end - remat.begin;
// Rewrite all operation indices.
for (auto& tensor : tensors_) {
std::for_each(std::lower_bound(tensor.operations.begin(),
tensor.operations.end(), remat.insert),
tensor.operations.end(),
[&](int& iop) { iop += num_inserted; });
}
operations_.insert(operations_.begin() + remat.insert, num_inserted, {});
// Copy all tensors dependencies of the old region to the new region.
// For any output tensor, a new tensor will be created.
std::vector<std::pair<int, int>> new_tensors;
for (int iop_old = remat.begin, iop_new = remat.insert; iop_old < remat.end;
++iop_old, ++iop_new) {
for (const auto itensor : operations_[iop_old].tensors) {
if (tensors_[itensor].first_use() == iop_old) {
new_tensors.emplace_back(itensor, AddTensor(tensors_[itensor].size));
}
AddUse(iop_new, itensor);
}
}
std::sort(new_tensors.begin(), new_tensors.end());
// Let all inserted + downstream operations refer to the new tensors.
for (int iop = remat.insert; iop < operations_.size(); ++iop) {
// We have to copy, otherwise we're modifying the set during the iteration.
for (const int old_tensor : std::vector<int>(operations_[iop].tensors)) {
const auto new_tensor =
std::lower_bound(new_tensors.begin(), new_tensors.end(),
std::make_pair(old_tensor, 0));
if (new_tensor != new_tensors.end() && new_tensor->first == old_tensor) {
DelUse(iop, old_tensor);
AddUse(iop, new_tensor->second);
}
}
}
}
void Rematerializer::RunGreedyAlgorithm(const int max_cost,
const int max_block_length,
const SizeT min_savings) {
const bool unlimited_cost = (max_cost < 0);
for (int min_block_length = 1, cost = 0;
min_block_length <= max_block_length &&
(unlimited_cost || cost <= max_cost);
min_block_length *= 2) {
while (unlimited_cost || cost <= max_cost) {
const auto [peak, remat] = FindBestRemat(
/*min_savings*/ min_savings,
/*begin_len=*/min_block_length,
/*end_len=*/
std::min(1 + (unlimited_cost
? max_block_length
: std::min(max_block_length, max_cost - cost)),
2 * min_block_length));
if (remat.begin == remat.end) break;
Remat(remat);
ApplyRemat(remat);
cost += (remat.end - remat.begin);
}
}
}
} // namespace TFL
} // namespace mlir
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer.h
+262
View File
@@ -0,0 +1,262 @@
/* 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_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
#define TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
// This file declares the Rematerializer class, which is used by an MLIR-based
// set of transformations for TFLite IR that lower memory usage by redoing
// operations with small inputs and large outputs instead of keeping the result
// in memory. This class allows us to compactly and efficiently represent the
// (idealized) memory profile of a TFLite graph and simulate the effect of
// re-inserting operations on that memory profile.
#include <algorithm>
#include <cinttypes>
#include <cstdint>
#include <tuple>
#include <vector>
namespace mlir {
namespace TFL {
// A class that
// (1) Encodes in concise form the memory requirements of a computational graph
// (2) Allows for the fast simulation of changes to the peak memory requirement
// under rematerialization of intermediate results in the graph
// (3) Implements a greedy algorithm for finding rematerializations of
// intermediate results in that graph to lower peak memory requirements.
class Rematerializer {
public:
Rematerializer() = default;
virtual ~Rematerializer() = default;
// The type used for memory sizes (in bytes) and differences thereof.
using SizeT = int64_t;
// The memory profile: The i-th element gives the amount of memory
// that is needed when performing the i-th operation. This is the
// sum of the sizes of
//
// (1) input tensors of that operation,
// (2) output tensors of that operation,
// (3) output tensors of preceding operations that are input tensors
// of subsequent operations.
using MemProfile = std::vector<SizeT>;
// Used for specifying memory consumption at a certain operation in the
// computational graph.
struct MemSpec {
int op_index; // The index of the operation
SizeT size; // The amount of memory needed in order to execute this
// operation, i.e., the sum of input and output sizes and the
// sizes of outputs of previous operations that are needed as
// inputs of subsequent operations.
explicit MemSpec(int op_index = 0, SizeT size = 0)
: op_index(op_index), size(size) {}
};
static bool BySize(const MemSpec& a, const MemSpec& b) {
return std::tie(a.size, a.op_index) < std::tie(b.size, b.op_index);
}
static bool ByOpIndex(const MemSpec& a, const MemSpec& b) {
return std::tie(a.op_index, a.size) < std::tie(b.op_index, b.size);
}
// Specifies an elementary rematerialization operation: The operations in
// operations [`begin`, `end`) will be rescheduled before operation `insert`.
// A valid `RematSpec` requires begin <= end <= insert <= number of
// operations. Note that (1) `end` is exclusive -- begin == end signifies a
// trivial RematSpec (no operation will be rescheduled), (2) the
// zero-initialized RematSpec {} is trivial and always valid.
struct RematSpec {
int begin;
int end;
int insert;
};
// Gives the peak memory location and size after inserting operations
// according to `remat` (but doesn't actually insert them.) Ties are broken
// towards later locations. `remat` must be valid (see above).
MemSpec GetPeakMemory(const RematSpec& remat = {}) const;
// Gives memory profile after inserting operations according to `remat` (but
// doesn't actually insert them). `remat` must be valid (see above).
MemProfile GetMemProfile(const RematSpec& remat = {}) const;
// Runs the greedy incremental block algorithm: Finds a sequence of
// rematerializations of block size up to max_block_length, each reducing peak
// memory by at least min_savings. If max_cost >= 0, at most max_cost
// operations will be re-inserted. For each rematerialization found,
// ApplyRemat is invoked (which can be used to apply the rematerialization to
// the higher- level representation, e.g., MLIR, flatbuffer, ...)
void RunGreedyAlgorithm(int max_cost, int max_block_length,
SizeT min_savings);
virtual void ApplyRemat(const RematSpec& remat) {}
protected:
// Rematerializes the outputs of the operations [`remat.begin`, `remat.end`)
// before operation remat.insert by copying that operation range before
// remat.insert and updating tensor references so that any operation that can
// will make use of the rematerialized outputs rather than the original ones.
// `remat` must be valid (see above).
void Remat(const RematSpec& remat);
// The protected methods below are to be used by derived classes to create the
// low-level (this class) representation from a high-level one.
// Creates a new tensor-like object that takes `size` bytes. Returns a
// contiguous increasing index for each new object, starting at 0.
int AddTensor(SizeT size);
// Creates an operation. If `is_stateful`, the operation (and any block of
// operations containing it) will never be considered for rematerialization.
// Returns a contiguous increasing index for each new object, starting at 0.
int AddOperation(bool is_stateful);
// The operator with index `ioperation` will be assumed to produce and/or
// consume the tensor with index `itensor`. NoOp if that's already the case.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void AddUse(int ioperation, int itensor);
// Undoes an AddUse(ioperation, itensor). NoOp if there was no prior `AddUse`.
// The arguments must be valid indices (i.e., obtained with
// `AddOperation`/`AddTensor`).
void DelUse(int ioperation, int itensor);
private:
// Find the best remat operation that saves at least `min_savings` bytes for a
// block of operators with a length is between [`begin_len`, `end_len`).
// 'Best' means with the highest savings, ties are broken towards shorter
// blocks.
std::tuple<SizeT, RematSpec> FindBestRemat(SizeT min_savings, int begin_len,
int end_len) const;
// Optimization: Estimate (from above) the remat savings of instruction block
// [begin, end) after operation `peak_location`
SizeT MaxSavings(int begin, int end, int peak_loc) const;
// If I want to remat ops [begin, end) after the op at operation `peak_loc`,
// find the latest point at which to reinsert them (the op before which to
// insert.)
int FindBestRematPoint(int begin, int end, int peak_loc) const;
// The memory objects.
struct Tensor {
SizeT size; // The size of the object (in bytes.)
std::vector<int> operations; // The operations it is used in. This vector
// is kept sorted + unique.
// The operation that makes the first use of this tensor.
int first_use() const { return *operations.begin(); }
// The operation that makes the last use of this tensor.
int last_use() const { return *operations.rbegin(); }
};
// The operators.
struct Operation {
bool is_stateful = false; // Results of an Operation can be rematerialized
// only if `!is_stateful`. This probably should
// be replaced with a more-fine grained
// approach--for example, the results of a "read
// resource variable" operation can be
// rematerialized as long as this doesn't happen
// after the corresponding "write resource
// variable" operation.
std::vector<int> tensors; // The tensors that are used (input or output) by
// this operation. They needn't correspond to
// tensors in the TF graph -- we may add fake
// tensors to model memory consumed in addition
// to input and output tensors. This vector is
// kept sorted + unique.
SizeT alloc = 0; // The number of bytes that need to be allocated before
// this operation.
SizeT dealloc = 0; // The number of bytes that can be deallocated after
// this operation.
};
// Given the current state of `operations_` and `tensors_`, return a vector of
// corrections that transform the current memory profile into the one that we
// would get after applying `remat`.
//
// The memory profile of a sequence of operations is the partial sum of the
// sizes of the allocations that are necessary before an operation and the
// negative sizes of the deallocations that are possible after the previous
// operation.
//
// If we modify the operation sequence by cloning an operation range, that
// memory profile will change--cloning makes it necessary to extend the
// lifetime of some tensors, while other tensors can be deallocated early and
// rematerialized later.
//
// This method represents these changes in compact form: It returns a vector
// of (position of operation, delta) pairs in lexicographic order; one
// obtains the memory profile after `remat` by adding the deltas from any
// entries (i, delta) to the i-th entry of the partial sum.
//
// This allows us to efficiently compute the change to the peak of a memory
// profile due to cloning an operation range without having to actually clone
// that range and without having to build a profile vector.
//
// The returned vector has at most 2 entries for each tensor referenced in
// [remat.begin, remat.end). There may be multiple entries for a single
// operation position; operation positions refer to the sequence *after*
// cloning [`remat.begin`, `remat.end`) before `remat.insert`.
std::vector<MemSpec> GetDeltas(const RematSpec& remat) const;
// Helper template: Iterates through all `MemSpec`s (i.e., operation
// index/memory usage pairs) for the current graph in operation order and
// calls `mapper` on them. This is an optimization -- by instantiating with an
// appropriate `Mapper`, it allows us to e.g. compute the peak memory without
// having to instantiate an actual memory profile vector.
template <class Mapper>
void MapMem(const Mapper& mapper, const RematSpec& remat) const {
const auto deltas = GetDeltas(remat);
const auto len = (remat.end - remat.begin);
auto idelta = deltas.begin();
for (MemSpec m; m.op_index < operations_.size() + len; ++m.op_index) {
// Are we in the cloned portion of the new operation sequence?
// Then all alloc/dealloc information must come from deltas.
const bool patch =
(m.op_index >= remat.insert) && (m.op_index < remat.insert + len);
// Are we past the insertion portion of the new operation sequence?
// Then we need to convert indices back to the original sequence.
const int shift = (m.op_index >= remat.insert + len) ? len : 0;
m.size += patch ? 0 : operations_[m.op_index - shift].alloc;
// deltas is sorted by location; apply any corrections to the current
// operator.
for (; idelta != deltas.end() && idelta->op_index == m.op_index;
++idelta) {
m.size += idelta->size;
}
mapper(m);
m.size -= patch ? 0 : operations_[m.op_index - shift].dealloc;
}
}
std::vector<Operation> operations_;
std::vector<Tensor> tensors_;
};
} // namespace TFL
} // namespace mlir
#endif // TENSORFLOW_COMPILER_MLIR_LITE_EXPERIMENTAL_REMAT_REMATERIALIZER_H_
tensorflow/compiler/mlir/lite/experimental/remat/rematerializer_test.cc
+519
View File
@@ -0,0 +1,519 @@
/* 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/compiler/mlir/lite/experimental/remat/rematerializer.h"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <initializer_list>
#include <random>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
namespace mlir {
namespace TFL {
namespace {
using ::testing::ElementsAre;
using ::testing::ElementsAreArray;
using ::testing::Eq;
using ::testing::FieldsAre;
using ::testing::StrictMock;
class RematTest : public ::testing::Test {
protected:
class TestableRematerializer : public Rematerializer {
public:
using Rematerializer::AddOperation;
using Rematerializer::AddTensor;
using Rematerializer::AddUse;
using Rematerializer::DelUse;
using Rematerializer::Remat;
};
TestableRematerializer r_;
};
TEST_F(RematTest, TensorUseSimple) {
for (int i = 0; i < 6; ++i) {
r_.AddOperation(/*is_stateful=*/false);
r_.AddTensor(/*size=*/1 << i);
}
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(4)));
r_.AddUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 4, 4, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/2, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 4, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(4), Eq(4)));
r_.DelUse(/*ioperation=*/4, /*itensor=*/2);
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(0, 0, 0, 0, 0, 0));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(5), Eq(0)));
}
TEST_F(RematTest, TensorUseMany) {
constexpr int n = 6;
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1 << (n - i - 1)));
}
for (int i = 0; i < n; ++i) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/n - 1 - i);
}
EXPECT_THAT(r_.GetMemProfile(), ElementsAreArray({32, 48, 56, 60, 62, 63, 63,
62, 60, 56, 48, 32}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(6), Eq(63)));
}
TEST_F(RematTest, PeakTiesAreBrokenInFavorOfLaterOperations) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
ASSERT_THAT(r_.GetMemProfile(), ElementsAreArray({100, 1, 100}));
EXPECT_THAT(r_.GetPeakMemory(), FieldsAre(Eq(2), Eq(100)));
}
TEST_F(RematTest, RematRecreatesOutput) {
r_.AddUse(r_.AddOperation(/*is_stateful=*/false), r_.AddTensor(100));
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// f2()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(100, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/0, /*end=*/1, /*insert=*/2}),
ElementsAre(100, 0, 100));
r_.Remat({/*begin=*/0, /*end=*/1, /*insert=*/2});
// /* after: */
// %0 = f1()
// f2()
// %1 = f1()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(100, 0, 100));
EXPECT_THAT(r_.AddTensor(0), 2);
}
TEST_F(RematTest, RematExtendsInputAndRecreatesOutput) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/1, 0);
r_.AddOperation(/*is_stateful=*/false);
r_.AddOperation(/*is_stateful=*/false);
// /* before: */
// %0 = f1()
// %1 = f2(%0)
// f3()
// f4()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 0, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/2, /*insert=*/3}),
ElementsAre(1, 101, 1, 101, 0));
r_.Remat({/*begin=*/1, /*end=*/2, /*insert=*/3});
// /* after: */
// %0 = f1()
// %1 = f2(%0)
// f3() /* %0 is kept alive */
// %2 = f2(%0)
// f4()
EXPECT_THAT(r_.GetMemProfile(), ElementsAre(1, 101, 1, 101, 0));
EXPECT_THAT(r_.AddTensor(0), 3);
}
TEST_F(RematTest, BlockRematDuplicatesIntraBlockValues) {
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(10));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(100));
r_.AddUse(/*ioperation=*/r_.AddOperation(/*is_stateful=*/false),
/*itensor=*/r_.AddTensor(1000));
r_.AddOperation(/*is_stateful=*/false);
r_.AddUse(/*ioperation=*/1, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/0);
r_.AddUse(/*ioperation=*/2, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/0);
r_.AddUse(/*ioperation=*/3, /*itensor=*/1);
r_.AddUse(/*ioperation=*/3, /*itensor=*/2);
// /* before */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
ASSERT_THAT(r_.GetMemProfile(), ElementsAre(1, 11, 111, 1111, 0));
EXPECT_THAT(r_.GetMemProfile({/*begin=*/1, /*end=*/4, /*insert=*/5}),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
r_.Remat({/*begin=*/1, /*end=*/4, /*insert=*/5});
EXPECT_THAT(r_.GetMemProfile(),
ElementsAre(1, 11, 111, 1111, 1, 11, 111, 1111));
EXPECT_THAT(r_.AddTensor(0), 7);
// /* after */
// %0 = f1()
// %1 = f2(%0)
// %2 = f3(%0,%1)
// %3 = f4(%0,%1,%2)
// f5()
// %4 = f2(%0)
// %5 = f3(%0,%4)
// %6 = f4(%0,%4,%5)
}
class RematSimulationTest : public testing::Test {
protected:
class RandomRemat : public Rematerializer {
public:
using Rematerializer::Remat; // For testing.
RandomRemat(const int num_operations, const int num_tensors,
const int num_uses, std::mt19937& rng) {
std::uniform_int_distribution<int> some_size_log(0, 16);
std::uniform_int_distribution<int> some_tensor(0, num_tensors - 1);
std::uniform_int_distribution<int> some_operation(0, num_operations - 1);
for (int i = 0; i < num_tensors; ++i) {
AddTensor(SizeT{1} << some_size_log(rng));
}
for (int i = 0; i < num_operations; ++i) {
AddOperation(/*is_stateful=*/false);
}
for (int i = 0; i < num_uses; ++i) {
AddUse(some_operation(rng), some_tensor(rng));
}
}
};
};
TEST_F(RematSimulationTest, SimulationAgreesWithReality) {
constexpr int kNumOperations = 128;
constexpr int kNumTensors = 32;
constexpr int kNumUses = kNumOperations * kNumTensors / 4;
std::mt19937 rng;
for (int i = 0; i < 1024; ++i) {
RandomRemat remat(kNumOperations, kNumTensors, kNumUses, rng);
// Worst-case scenario: we might double the length of the computation
// schedule each time we remat, so only a few iterations...
std::array<int, 3> randos;
const auto& [begin, end, insert] = randos;
for (int i = 0, num_operations = kNumOperations; i < 4;
++i, num_operations += end - begin) {
std::uniform_int_distribution<int> some_op(0, num_operations - 1);
for (auto& rando : randos) {
rando = some_op(rng);
}
// We need begin <= end <= insert.
std::sort(randos.begin(), randos.end());
const Rematerializer::RematSpec spec{begin, end, insert};
const auto simulated_profile = remat.GetMemProfile(spec);
remat.Remat(spec);
const auto actual_profile = remat.GetMemProfile();
EXPECT_THAT(simulated_profile, ElementsAreArray(actual_profile));
}
}
}
class GreedyRematTest : public testing::Test {
protected:
// A Rematerializer for a graph whose memory profile is a sequence of
// "rainbows" due to nested ops (similar to what happens in backwards pass
// gradient calculations.)
class RainbowRemat : public Rematerializer {
public:
// `sizes` is a vector of vector of sizes. Each sizes vector of length n
// generates a `rainbow` profile of the SSA form
//
// %1 = f1()
// %2 = f2()
// ...
// %n = fn()
// gn(%n)
// ...
// g2(%2)
// g1(%1)
//
// with the sizes of intermediates %1,...%n given by the sizes entries.
//
// If extra_ops > 0, each assignment above becomes a sequence
// %i.1 = fi.1()
// %i.2 = fi.1(%i.1)
// %i.3 = fi.1(%i.2)
// ...
// %i = f1(%i.m)
// where the 'dotted' intermediates have size `extra_size`. This
// allows for testing the effect of window sizes.
explicit RainbowRemat(const std::vector<std::vector<int>>& sizes,
int extra_ops = 0, SizeT extra_size = 0) {
for (const auto& rainbow : sizes) {
int tensor = 0;
int op = 0;
for (const auto& size : rainbow) {
for (int i = 0; i < extra_ops; ++i) {
op = AddOperation(/*is_stateful=*/false);
if (i != 0) {
AddUse(op, tensor);
}
tensor = AddTensor(extra_size);
AddUse(op, tensor);
}
// using negative sizes to signal forbidden operations.
op = AddOperation(/*is_stateful=*/size < 0);
if (extra_ops > 0) {
AddUse(op, tensor);
}
tensor = AddTensor(std::abs(size));
AddUse(op, tensor);
}
for (int i = 0; i < rainbow.size(); ++i) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, tensor - i);
}
}
}
};
// Similar to above: A multilayer perceptron. The forward pass is
// f[n](f[n-1](...(f[1]())), with dimensions |f[1]|...|f[n]| handed in the
// constructor. The backward pass calculates the loss gradient for a single
// weight in the first layer.
class MlpRemat : public Rematerializer {
public:
explicit MlpRemat(const std::vector<int>& sizes) {
int forward_tensor = -1;
int backward_tensor = -1;
int op = -1;
// Forward pass:
for (const int size : sizes) {
op = AddOperation(/*is_stateful=*/false);
if (forward_tensor >= 0) AddUse(op, forward_tensor);
forward_tensor = AddTensor(size);
AddUse(op, forward_tensor);
}
// Backward pass: Right-multiply the jacobian from the outside in.
// dLoss/df[n] * df[n]/df[n-1] * ...
// The i-th term g[i] depends on g[i-1] and f[n-i] and has size |f[n-i]|.
for (; forward_tensor >= 0; --forward_tensor) {
op = AddOperation(/*is_stateful=*/false);
AddUse(op, forward_tensor);
if (backward_tensor >= 0) AddUse(op, backward_tensor);
backward_tensor = AddTensor(sizes[forward_tensor]);
AddUse(op, backward_tensor);
}
}
// We will also instrument the actual optimizations performed.
MOCK_METHOD(void, ApplyRemat, (const RematSpec&));
};
};
TEST_F(GreedyRematTest, MlpBasic) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 1, 1}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) l i o 3
// %3 = g2( %2) l l i o 4
// %4 = g1(%3,%1) l i i o 4
// %5 = g0(%4,%0) i i o 3
ASSERT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 3, 4, 4, 3}));
// Little we can do here -- remat %0 before %5
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/5)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 2
// %2 = f2(%1) i o 2
// %3 = g2( %2) l i o 3
// %4 = g1(%3,%1) i i o 3
// %0' = f0() o l 2
// %5 = g0(%4,%0') i i o 3
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(), ElementsAreArray({1, 2, 2, 3, 3, 2, 3}));
}
TEST_F(GreedyRematTest, MlpBinary) {
StrictMock<MlpRemat> remat(std::vector<int>({1, 2, 4, 8}));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) l i o 7
// %3 = f3(%2) l l i o 15
// %4 = g3( %3) l l l i o 23
// %5 = g2(%4 %2) l l i i o 19
// %6 = g1(%5,%1) l i i o 9
// %7 = g0(%6,%0) i i o 4
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 19, 9, 4}));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/2,
/*end=*/3,
/*insert=*/5)));
EXPECT_CALL(remat, ApplyRemat(FieldsAre(/*begin=*/0,
/*end=*/1,
/*insert=*/8)));
// (o)ut, (i)n, (l)ive: 0 1 2 3 4 5 6 7 Sum of live sizes
// %0 = f0() o 1
// %1 = f1(%0) i o 3
// %2 = f2(%1) i o 6
// %3 = f3(%2) l i o 14
// %4 = g3( %3) l i o 18
// %2'= f2(%1) l o l 14
// %5 = g2(%4 %2) l i i o 18
// %6 = g1(%5,%1) i i o 8
// %0'= f(0) o l 3
// %7 = g0(%6,%0) i i o 4
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 6, 14, 18, 14, 18, 8, 3, 4}));
}
TEST_F(GreedyRematTest, SimpleMax) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleMaxLongWindow) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/4,
/*min_savings=*/1);
// Profile is flattened to its minimum--16.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, SimpleSizeThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Only do remats of at least 4 bytes of savings -- this will lower the
// profile by 4 + 8 instead of 1 + 2 + 4 + 8 as before.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/4);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 11, 19, 19, 11, 11, 7, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleCostThreshold) {
RainbowRemat remat({{1, 2, 4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Do at most 1 remat -- this will lower the profile by 8.
remat.RunGreedyAlgorithm(/*max_cost=*/1, /*max_block_length=*/1,
/*min_savings=*/1);
// Only as single remat is done -- this will be the best one possible,
// lowering the profile by 8 instead of the maximum 1 + 2 + 4 + 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1}));
}
TEST_F(GreedyRematTest, SimpleForbiddenOps) {
// Operator generating size-4 tensor is stateful, so it won't be materialized.
RainbowRemat remat({{1, 2, -4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 31, 31, 15, 7, 3, 1}));
// Best we can do is lower the profile to 8 + 2 + 1.
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 12, 20, 20, 12, 12, 4, 2, 2, 1, 1}));
}
TEST_F(GreedyRematTest, DoubleMax) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/1,
/*min_savings=*/1);
// Two rainbows; the first is lowered by 1 + 2 + 4 + 8 == 15, the second by 4
// + 8 == 12.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 2, 4, 8, 16, 16, 8, 8, 4, 4, 2,
2, 1, 1, 4, 8, 16, 16, 8, 8, 4, 4}));
}
TEST_F(GreedyRematTest, DoubleCostThreshold) {
// Generate a profile with two local maxima of size 31 and 28.
RainbowRemat remat({{1, 2, 4, 8, 16}, {4, 8, 16}});
ASSERT_THAT(remat.GetMemProfile(),
ElementsAreArray(
{1, 3, 7, 15, 31, 31, 15, 7, 3, 1, 4, 12, 28, 28, 12, 4}));
remat.RunGreedyAlgorithm(/*max_cost=*/2, /*max_block_length=*/1,
/*min_savings=*/1);
// Profile can be flattened twice--first, the global maximum of 31 is reduced
// by 8; then the newly-global maximum of 28 is reduced by 8.
EXPECT_THAT(remat.GetMemProfile(),
ElementsAreArray({1, 3, 7, 15, 23, 23, 15, 15, 7, 3, 1, 4, 12, 20,
20, 12, 12, 4}));
}
TEST_F(GreedyRematTest, SingleLongerBlocksByWindowSize) {
std::vector<Rematerializer::SizeT> best_for_window_size;
for (int window_size : {0, 1, 2, 3, 4, 5}) {
RainbowRemat remat({{1, 2, 4, 8}}, 2, 16);
remat.RunGreedyAlgorithm(/*max_cost=*/-1, /*max_block_length=*/window_size,
/*min_savings=*/1);
best_for_window_size.push_back(remat.GetPeakMemory().size);
}
EXPECT_THAT(best_for_window_size, ElementsAreArray({44, 36, 36, 32, 32, 32}));
}
} // namespace
} // namespace TFL
} // namespace mlir