wehub-resource-sync
316 lines
12 KiB
C++
316 lines
12 KiB
C++
/* 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_DTENSOR_CC_SLICE_UTIL_H_
|
|
#define TENSORFLOW_DTENSOR_CC_SLICE_UTIL_H_
|
|
|
|
#include <optional>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
|
|
#include "tensorflow/core/platform/statusor.h"
|
|
#include "tensorflow/dtensor/cc/tensor_layout.h"
|
|
|
|
namespace tensorflow {
|
|
namespace dtensor {
|
|
namespace slice_util {
|
|
|
|
// Defines a token of the strided slicing mini-language.
|
|
// Refer to the definition of StridedSlice Op for the informal definition of
|
|
// the language. During slicing, axes of the input tensor are processed one
|
|
// by one according to the tokens of the slicing spec vector.
|
|
struct Token {
|
|
enum TokenType {
|
|
REGULAR, // Slice the current axis by begin/end/begin_mask/end_mask and
|
|
// stride.
|
|
NEW_AXIS, // Add a new axis at the current location to the output.
|
|
ELLIPSIS, // Copy over following axes to the output till the ellipsis ends.
|
|
SHRINK_AXIS // Like a regular axis, but sequeeze this axis from output
|
|
// after slicing.
|
|
} token_type;
|
|
|
|
int64_t begin = 0; // Begin of the slice.
|
|
int64_t end = 0; // End of the slice.
|
|
int64_t stride = 0; // Stride of the slice.
|
|
bool dynamic_mask = false; // If begin, end, or stride is a dynamic value.
|
|
bool begin_mask = false; // True if the begin is maximal.
|
|
bool end_mask = false; // True if the end is maximal.
|
|
|
|
Token() = default;
|
|
Token(TokenType token_type, int64_t begin, int64_t end, int64_t stride,
|
|
bool dynamic_mask = false, bool begin_mask = false,
|
|
bool end_mask = false)
|
|
: token_type(token_type),
|
|
begin(begin),
|
|
end(end),
|
|
stride(stride),
|
|
dynamic_mask(dynamic_mask),
|
|
begin_mask(begin_mask),
|
|
end_mask(end_mask) {}
|
|
|
|
// Normalizes the token such that (end - begin) is evenly divided by stride,
|
|
// and the result equals the total elements after the slicing.
|
|
Token normalize(int64_t dim_size) const;
|
|
std::optional<Token> GetLocalToken(int64_t dim_size,
|
|
int64_t num_shards) const;
|
|
};
|
|
|
|
// TODO(feyu): is there a C++ way to do vari args and templates move this out
|
|
// of this class?
|
|
template <typename T, typename... Types>
|
|
StatusOr<T> CreateAndRun(const std::vector<Token>& tokens, Types... args) {
|
|
T visitor(args...);
|
|
TF_RETURN_IF_ERROR(visitor.Run(tokens));
|
|
return visitor;
|
|
}
|
|
|
|
class TokenProcessor {
|
|
public:
|
|
explicit TokenProcessor(int64_t input_rank) : input_rank_(input_rank) {}
|
|
virtual ~TokenProcessor() = default;
|
|
|
|
absl::Status Run(const std::vector<Token>& tokens);
|
|
|
|
protected:
|
|
// Loop for an ellipsis or the unconsumed axes in the end.
|
|
bool VisitLoop(int64_t input_rank, int64_t output_rank, int64_t ellipsis_size,
|
|
int64_t* input_index, int64_t* output_index);
|
|
|
|
virtual void VisitImplicitAxis(int64_t input_index, int64_t output_index) = 0;
|
|
|
|
virtual void VisitEllipsisAxis(const Token& token) = 0;
|
|
|
|
virtual void VisitShrinkAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) = 0;
|
|
|
|
virtual void VisitNewAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) = 0;
|
|
|
|
virtual void VisitRegularAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) = 0;
|
|
|
|
virtual void PrepareResults(int64_t spec_rank, int64_t input_rank,
|
|
int64_t output_rank) = 0;
|
|
|
|
virtual absl::Status FinalizeResults(int64_t input_rank,
|
|
int64_t output_rank) = 0;
|
|
|
|
private:
|
|
const int64_t input_rank_;
|
|
};
|
|
|
|
// Forward layout inference of from a StridedSlice token vector.
|
|
//
|
|
// For value_layout = StridedSlice(input_layout, tokens)
|
|
//
|
|
// The inference consumes input_layout, and produces:
|
|
// - a planned expander_input_layout that is suitable for SPMD expansion.
|
|
// - a planned expander_value_layout that is suitable for SPMD expansion.
|
|
// - a local_tokens vector for the arguments of the post-SPMD StridedSliceOp.
|
|
// expander_input_layout and expander_value_layout are consistent with
|
|
// local_tokens.
|
|
class ForwardLayoutInference : public TokenProcessor {
|
|
public:
|
|
ForwardLayoutInference(const Layout& input_layout,
|
|
const llvm::ArrayRef<int64_t> input_shape)
|
|
: TokenProcessor(input_shape.size()),
|
|
input_layout_(input_layout),
|
|
input_shape_(input_shape),
|
|
input_sharding_(input_layout.sharding_spec_strs()) {}
|
|
|
|
const Layout& expander_value_layout() const { return expander_value_layout_; }
|
|
|
|
const Layout& expander_input_layout() const { return expander_input_layout_; }
|
|
|
|
const std::vector<Token>& local_tokens() const { return local_tokens_; }
|
|
|
|
protected:
|
|
void VisitEllipsisAxis(const Token& token) override {
|
|
local_tokens_.push_back(token);
|
|
}
|
|
|
|
void VisitImplicitAxis(int64_t input_index, int64_t output_index) override {
|
|
expander_input_sharding_.push_back(input_sharding_[output_index]);
|
|
expander_value_sharding_.push_back(input_sharding_[output_index]);
|
|
}
|
|
|
|
void VisitShrinkAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
local_tokens_.push_back(token);
|
|
expander_input_sharding_.push_back(Layout::kUnshardedDim);
|
|
// Skips this axis from values, since it will be removed from the inputs.
|
|
}
|
|
|
|
void VisitNewAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
local_tokens_.push_back(token);
|
|
expander_value_sharding_.push_back(Layout::kUnshardedDim);
|
|
}
|
|
|
|
void VisitRegularAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
auto local_token = token.GetLocalToken(
|
|
/*dim_size=*/input_shape_[input_index],
|
|
/*num_shards*/ input_layout_.num_shards_for_dim(input_index));
|
|
std::string sharding = input_sharding_[input_index];
|
|
if (local_token.has_value()) {
|
|
local_tokens_.push_back(*local_token);
|
|
} else {
|
|
sharding = Layout::kUnshardedDim;
|
|
local_tokens_.push_back(token);
|
|
}
|
|
expander_value_sharding_.push_back(sharding);
|
|
expander_input_sharding_.push_back(sharding);
|
|
}
|
|
|
|
void PrepareResults(int64_t spec_rank, int64_t input_rank,
|
|
int64_t output_rank) override {
|
|
local_tokens_.reserve(spec_rank);
|
|
expander_input_sharding_.reserve(input_rank);
|
|
expander_value_sharding_.reserve(output_rank);
|
|
}
|
|
|
|
absl::Status FinalizeResults(int64_t input_rank,
|
|
int64_t output_rank) override {
|
|
DCHECK_EQ(expander_input_sharding_.size(), input_rank);
|
|
DCHECK_EQ(expander_value_sharding_.size(), output_rank);
|
|
TF_ASSIGN_OR_RETURN(
|
|
expander_input_layout_,
|
|
Layout::GetLayout(expander_input_sharding_, input_layout_.mesh()));
|
|
TF_ASSIGN_OR_RETURN(
|
|
expander_value_layout_,
|
|
Layout::GetLayout(expander_value_sharding_, input_layout_.mesh()));
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
private:
|
|
const Layout& input_layout_;
|
|
const llvm::ArrayRef<int64_t> input_shape_;
|
|
std::vector<std::string> input_sharding_;
|
|
std::vector<std::string> expander_value_sharding_;
|
|
std::vector<std::string> expander_input_sharding_;
|
|
// Outputs
|
|
Layout expander_value_layout_;
|
|
Layout expander_input_layout_;
|
|
std::vector<Token> local_tokens_;
|
|
};
|
|
|
|
// Backward layout inference for a StridedSlice token vector.
|
|
//
|
|
// For value_layout = StridedSlice(input_layout, tokens)
|
|
//
|
|
// The inference consumes value_layout, and produces:
|
|
// - a planned expander_input_layout that is suitable for SPMD expansion.
|
|
// - a planned expander_value_layout that is suitable for SPMD expansion.
|
|
// - a local_tokens vector for the arguments of the post-SPMD StridedSliceOp.
|
|
// expander_input_layout and expander_value_layout are consistent with
|
|
// local_tokens.
|
|
class BackwardLayoutInference : public TokenProcessor {
|
|
public:
|
|
BackwardLayoutInference(const Layout& value_layout,
|
|
const llvm::ArrayRef<int64_t> input_shape)
|
|
: TokenProcessor(input_shape.size()),
|
|
value_layout_(value_layout),
|
|
input_shape_(input_shape),
|
|
value_sharding_(value_layout.sharding_spec_strs()) {}
|
|
|
|
const Layout& expander_input_layout() const { return expander_input_layout_; }
|
|
|
|
const Layout& expander_value_layout() const { return expander_value_layout_; }
|
|
|
|
const std::vector<Token>& local_tokens() const { return local_tokens_; }
|
|
|
|
protected:
|
|
void VisitEllipsisAxis(const Token& token) override {
|
|
local_tokens_.push_back(token);
|
|
}
|
|
|
|
void VisitImplicitAxis(int64_t input_index, int64_t output_index) override {
|
|
expander_input_sharding_.push_back(value_sharding_[output_index]);
|
|
expander_value_sharding_.push_back(value_sharding_[output_index]);
|
|
}
|
|
|
|
void VisitShrinkAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
local_tokens_.push_back(token);
|
|
// There is no constraint on the input sharding, but we prefer to keep it
|
|
// unsharded to avoid inserting relayout toward the internal input layout.
|
|
expander_input_sharding_.push_back(Layout::kUnshardedDim);
|
|
}
|
|
|
|
void VisitNewAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
local_tokens_.push_back(token);
|
|
// No corresponding input axis.
|
|
expander_value_sharding_.push_back(Layout::kUnshardedDim);
|
|
}
|
|
|
|
void VisitRegularAxis(const Token& token, int64_t input_index,
|
|
int64_t output_index) override {
|
|
auto local_token = token.GetLocalToken(
|
|
/*dim_size=*/input_shape_[input_index],
|
|
/*num_shards*/ value_layout_.num_shards_for_dim(output_index));
|
|
if (local_token.has_value()) {
|
|
std::string sharding = value_sharding_[output_index];
|
|
local_tokens_.push_back(*local_token);
|
|
expander_input_sharding_.push_back(sharding);
|
|
expander_value_sharding_.push_back(sharding);
|
|
} else {
|
|
local_tokens_.push_back(token);
|
|
// There is no constraint on the input sharding, but we prefer to keep it
|
|
// unsharded to avoid inserting relayout toward the internal input layout.
|
|
expander_input_sharding_.push_back(Layout::kUnshardedDim);
|
|
expander_value_sharding_.push_back(Layout::kUnshardedDim);
|
|
}
|
|
}
|
|
|
|
void PrepareResults(int64_t spec_rank, int64_t input_rank,
|
|
int64_t output_rank) override {
|
|
local_tokens_.reserve(spec_rank);
|
|
expander_input_sharding_.reserve(input_rank);
|
|
expander_value_sharding_.reserve(output_rank);
|
|
}
|
|
|
|
absl::Status FinalizeResults(int64_t input_rank,
|
|
int64_t output_rank) override {
|
|
DCHECK_EQ(expander_input_sharding_.size(), input_rank);
|
|
DCHECK_EQ(expander_value_sharding_.size(), output_rank);
|
|
TF_ASSIGN_OR_RETURN(
|
|
expander_input_layout_,
|
|
Layout::GetLayout(expander_input_sharding_, value_layout_.mesh()));
|
|
TF_ASSIGN_OR_RETURN(
|
|
expander_value_layout_,
|
|
Layout::GetLayout(expander_value_sharding_, value_layout_.mesh()));
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
private:
|
|
const Layout& value_layout_;
|
|
const llvm::ArrayRef<int64_t> input_shape_;
|
|
std::vector<std::string> value_sharding_;
|
|
std::vector<std::string> expander_input_sharding_;
|
|
std::vector<std::string> expander_value_sharding_;
|
|
// Outputs
|
|
Layout expander_input_layout_;
|
|
Layout expander_value_layout_;
|
|
std::vector<Token> local_tokens_;
|
|
};
|
|
|
|
} // namespace slice_util
|
|
} // namespace dtensor
|
|
} // namespace tensorflow
|
|
#endif // TENSORFLOW_DTENSOR_CC_SLICE_UTIL_H_
|