experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #43226 from kaixih:reduce_ops_layout

Browse Source 
PiperOrigin-RevId: 333786383
Change-Id: Ifefb0a3d23cf7779858e7c011fd4195024ab9dc5
This commit is contained in:
TensorFlower Gardener 2020-09-25 12:56:20 -07:00
commit ab55c62645
7 changed files with 168 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/compiler/tests/data_format_ops_test.py
View File

@ -63,6 +63,22 @@ class XlaDataFormatDimMapTest(xla_test.XLATestCase):
self._test([-4, -3, -2, -1, 0, 1, 2, 3], "qwer", "rewq",
[3, 2, 1, 0, 3, 2, 1, 0])
self._test(0, "NDHWC", "NCDHW", 0)
self._test(1, "NDHWC", "NCDHW", 2)
self._test(2, "NDHWC", "NCDHW", 3)
self._test(3, "NDHWC", "NCDHW", 4)
self._test(4, "NDHWC", "NCDHW", 1)
self._test([1, 4], "NDHWC", "NCDHW", [2, 1])
self._test([1, 4, -2], "NDHWC", "NCDHW", [2, 1, 4])
self._test([1, -3, -2], "NDHWC", "NCDHW", [2, 3, 4])
self._test([[1, -4], [1, -1]], "NDHWC", "NCDHW", [[2, 2], [2, 1]])
self._test([1, -3, -2], "NDHWC", "NCDHW", [2, 3, 4])
self._test([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4], "NDHWC", "DHWNC",
[3, 0, 1, 2, 4, 3, 0, 1, 2, 4])
self._test([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4], "NDHWC", "WHDCN",
[4, 2, 1, 0, 3, 4, 2, 1, 0, 3])
class XlaPermuteOpTest(xla_test.XLATestCase):

23
tensorflow/compiler/tf2xla/kernels/data_format_ops.cc
View File

@ -35,15 +35,19 @@ class DataFormatDimMapOp : public XlaOpKernel {
OP_REQUIRES_OK(context, context->GetAttr("src_format", &src_format));
string dst_format;
OP_REQUIRES_OK(context, context->GetAttr("dst_format", &dst_format));
OP_REQUIRES(context, src_format.size() == 4,
errors::InvalidArgument(absl::StrCat(
"Source format must of length 4, received src_format = ",
src_format)));
OP_REQUIRES(context, src_format.size() == 4 or src_format.size() == 5,
errors::InvalidArgument(
absl::StrCat("Source format must of length 4 or 5, "
"received src_format = ",
src_format)));
OP_REQUIRES(
context, dst_format.size() == 4,
context, dst_format.size() == 4 or dst_format.size() == 5,
errors::InvalidArgument(absl::StrCat(
"Destination format must of length 4, received dst_format = ",
"Destination format must of length 4 or 5, received dst_format = ",
dst_format)));
for (int i = 0; i < src_format.size(); ++i) {
dst_idx_.push_back(-1);
}
for (int i = 0; i < src_format.size(); ++i) {
for (int j = 0; j < dst_format.size(); ++j) {
if (dst_format[j] == src_format[i]) {
@ -61,9 +65,10 @@ class DataFormatDimMapOp : public XlaOpKernel {
auto builder = context->builder();
xla::XlaOp dst_indices =
xla::ConstantR1(builder, absl::Span<const int32>(dst_idx_));
xla::XlaOp four = xla::ConstantR0<int32>(builder, 4);
const int dims = dst_idx_.size();
xla::XlaOp rank = xla::ConstantR0<int32>(builder, dims);
xla::XlaOp src_indices =
(xla::ConvertElementType(context->Input(0), xla::S32) + four) % four;
(xla::ConvertElementType(context->Input(0), xla::S32) + rank) % rank;
xla::XlaOp output =
xla::TorchIndexSelect(dst_indices, src_indices, /*dim=*/0);
context->SetOutput(
@ -71,7 +76,7 @@ class DataFormatDimMapOp : public XlaOpKernel {
}
private:
std::array<int32, 4> dst_idx_ = {{-1, -1, -1, -1}};
std::vector<int32> dst_idx_;
TF_DISALLOW_COPY_AND_ASSIGN(DataFormatDimMapOp);
};

31
tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer.cc
View File

@ -1371,11 +1371,35 @@ bool ReduceTransposer::IsReduceAxisSupported(
Status ReduceTransposer::TransposeNode(TransposeContext* context,
utils::MutableNodeView* node) {
DCHECK(IsReduceOp(*node->node()));
if (!ShouldProcess(*context, *node) || !IsFaninPortRankN(*node, 0, 4) ||
const auto& regular_fanin = node->GetRegularFanin(0);
const auto* output_shape_attr =
regular_fanin.node_view()->GetAttr(kAttrOutputShape);
const auto& shape = output_shape_attr->list().shape(0);
const int rank = shape.dim_size();
std::string src_format = context->src_format;
std::string dst_format = context->dst_format;
// Update the format from 4D to 5D layout if necessary.
bool allow_5d = rank == 5 && (src_format == "NHWC" || src_format == "NCHW") &&
(dst_format == "NHWC" || dst_format == "NCHW");
if (allow_5d) {
std::string src_format_3d = src_format == "NHWC" ? "NDHWC" : "NCDHW";
std::string dst_format_3d = dst_format == "NHWC" ? "NDHWC" : "NCDHW";
context->AssignDeviceAndDataFormats(context->target_device, src_format_3d,
dst_format_3d);
}
if (!ShouldProcess(*context, *node) || !IsFaninPortRankN(*node, 0, rank) ||
!IsReduceAxisSupported(*context, *node) ||
!IsAfterDstToSrcTransform(*context, *node)) {
// Change back to the original layout due to early exit.
if (allow_5d) {
context->AssignDeviceAndDataFormats(context->target_device, src_format,
dst_format);
}
return Status::OK();
}
VLOG(3) << "GenericLayoutOptimizer: transforming node '" << node->GetName()
<< "' with op '" << node->GetOp() << "' from data format '"
<< context->src_format << "' to '" << context->dst_format << "'";
TF_RETURN_IF_ERROR(UpdateFaninEdgesWithOp(context, {0}, node, kOpTranspose));
TF_RETURN_IF_ERROR(
UpdateFaninEdgesWithOp(context, {1}, node, kOpDataFormatDimMap));
@ -1383,6 +1407,11 @@ Status ReduceTransposer::TransposeNode(TransposeContext* context,
TF_RETURN_IF_ERROR(
UpdateFanoutEdgesWithOp(context, {0}, node, kOpTranspose));
}
// Change back the format from 5D to 4D layout.
if (allow_5d) {
context->AssignDeviceAndDataFormats(context->target_device, src_format,
dst_format);
}
return context->graph_view->GetMutationBuilder()->Apply();
}

9
tensorflow/core/kernels/data_format_ops.cc
View File

@ -37,14 +37,15 @@ class DataFormatDimMapOp : public OpKernel {
OP_REQUIRES_OK(context, context->GetAttr("src_format", &src_format));
string dst_format;
OP_REQUIRES_OK(context, context->GetAttr("dst_format", &dst_format));
OP_REQUIRES(context, src_format.size() == 4,
OP_REQUIRES(context, src_format.size() == 4 || src_format.size() == 5,
errors::InvalidArgument(strings::StrCat(
"Source format must of length 4, received src_format = ",
"Source format must of length 4 or 5, received "
"src_format = ",
src_format)));
OP_REQUIRES(
context, dst_format.size() == 4,
context, dst_format.size() == 4 || dst_format.size() == 5,
errors::InvalidArgument(strings::StrCat(
"Destination format must of length 4, received dst_format = ",
"Destination format must of length 4 or 5, received dst_format = ",
dst_format)));
dst_idx_ = Tensor(DT_INT32, {static_cast<int64>(src_format.size())});
for (int i = 0; i < src_format.size(); ++i) {

54
tensorflow/core/kernels/data_format_ops.h
View File

@ -28,24 +28,50 @@ template <typename Device, typename T>
struct DataFormatDimMap {
void operator()(const Device& d, typename TTypes<T>::ConstFlat x,
typename TTypes<T>::Flat y, const TTypes<int>::Vec dst) {
auto zero = x.constant(0);
auto one = x.constant(1);
auto two = x.constant(2);
if (dst.size() == 4) {
auto zero = x.constant(0);
auto one = x.constant(1);
auto two = x.constant(2);
auto f_zero = x.constant(dst(0));
auto f_one = x.constant(dst(1));
auto f_two = x.constant(dst(2));
auto f_three = x.constant(dst(3));
auto f_zero = x.constant(dst(0));
auto f_one = x.constant(dst(1));
auto f_two = x.constant(dst(2));
auto f_three = x.constant(dst(3));
auto four = x.constant(4);
auto x_mod = (x + four) % 4;
auto four = x.constant(4);
auto x_mod = (x + four) % 4;
auto is_zero = (x_mod == zero);
auto is_one = (x_mod == one);
auto is_two = (x_mod == two);
auto is_zero = (x_mod == zero);
auto is_one = (x_mod == one);
auto is_two = (x_mod == two);
y.device(d) = is_zero.select(
f_zero, is_one.select(f_one, is_two.select(f_two, f_three)));
y.device(d) = is_zero.select(
f_zero, is_one.select(f_one, is_two.select(f_two, f_three)));
} else {
auto zero = x.constant(0);
auto one = x.constant(1);
auto two = x.constant(2);
auto three = x.constant(3);
auto f_zero = x.constant(dst(0));
auto f_one = x.constant(dst(1));
auto f_two = x.constant(dst(2));
auto f_three = x.constant(dst(3));
auto f_four = x.constant(dst(4));
auto five = x.constant(5);
auto x_mod = (x + five) % 5;
auto is_zero = (x_mod == zero);
auto is_one = (x_mod == one);
auto is_two = (x_mod == two);
auto is_three = (x_mod == three);
y.device(d) = is_zero.select(
f_zero,
is_one.select(
f_one, is_two.select(f_two, is_three.select(f_three, f_four))));
}
}
};

36
tensorflow/python/grappler/layout_optimizer_test.py
View File

@ -221,6 +221,9 @@ class LayoutOptimizerTest(test.TestCase):
def _assert_map_nhwc_to_nchw(self, name, nodes):
self.assertIn(name + '-DimMapNHWCToNCHW-LayoutOptimizer', nodes)
def _assert_map_ndhwc_to_ncdhw(self, name, nodes):
self.assertIn(name + '-DataFormatDimMapNDHWCToNCDHW-LayoutOptimizer', nodes)
def _assert_vec_nchw_to_nhwc(self, name, nodes):
self.assertIn(name + '-VecPermuteNCHWToNHWC-LayoutOptimizer', nodes)
@ -1194,6 +1197,39 @@ class LayoutOptimizerTest(test.TestCase):
self._assert_trans_nchw_to_nhwc('LeakyReluGrad-0-0', nodes)
self.assertAllClose(output_val_ref, output_val, atol=1e-3)
@test_util.deprecated_graph_mode_only
def testReduceOpsFor5DTensors(self):
if test.is_gpu_available(cuda_only=True):
random_seed.set_random_seed(0)
x = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
w = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
conv3d = gen_nn_ops.conv3d(x, w, [1, 1, 1, 1, 1], 'SAME')
y = math_ops.reduce_mean(conv3d, [0, 1, 2, 3], keepdims=True)
output = array_ops.identity(y)
with session.Session(config=_get_config(False)) as sess:
output_val_ref = sess.run(output)
with session.Session(config=_get_config()) as sess:
metadata = config_pb2.RunMetadata()
output_val = sess.run(output, run_metadata=metadata)
nodes = []
num_transposes = 0
for node in metadata.cost_graph.node:
if _is_transpose(node.name):
num_transposes += 1
nodes.append(node.name)
# The reduce op Mean needs to dim map the input reduce index to NCDHW.
# Then, the output needs to be tranposed back to NDHWC.
expected_num_transposes = 2
self.assertEqual(expected_num_transposes, num_transposes)
self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
self._assert_map_ndhwc_to_ncdhw('Mean-1', nodes)
self._assert_trans_ncdhw_to_ndhwc('Mean-0-0', nodes)
self.assertAllClose(output_val_ref, output_val, atol=1e-3)
@test_util.deprecated_graph_mode_only
def testConv3D(self):
if test.is_gpu_available(cuda_only=True):

27
tensorflow/python/ops/nn_test.py
View File

@ -1235,6 +1235,33 @@ class DataFormatDimMapTest(test_lib.TestCase):
y_val = self.evaluate(y)
self.assertAllEqual(y_val, y_val_expected)
def testNDHWCtoNCDHW(self):
x_val = [1, -4, -3, -2]
y_val_expected = [2, 2, 3, 4]
x = constant_op.constant(x_val)
y = nn_ops.data_format_dim_map(x, src_format="NDHWC", dst_format="NCDHW")
with test_util.use_gpu():
y_val = self.evaluate(y)
self.assertAllEqual(y_val, y_val_expected)
def testNDHWCtoDHWNC(self):
x_val = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
y_val_expected = [3, 0, 1, 2, 4, 3, 0, 1, 2, 4]
x = constant_op.constant(x_val)
y = nn_ops.data_format_dim_map(x, src_format="NDHWC", dst_format="DHWNC")
with test_util.use_gpu():
y_val = self.evaluate(y)
self.assertAllEqual(y_val, y_val_expected)
def testDNHWCtoWHDCN(self):
x_val = [-5, -4, -3, -2, -1, 0, 1, 2, 3, 4]
y_val_expected = [4, 2, 1, 0, 3, 4, 2, 1, 0, 3]
x = constant_op.constant(x_val)
y = nn_ops.data_format_dim_map(x, src_format="NDHWC", dst_format="WHDCN")
with test_util.use_gpu():
y_val = self.evaluate(y)
self.assertAllEqual(y_val, y_val_expected)
def testArbitraryASCII(self):
x_val = [-4, -3, -2, -1, 0, 1, 2, 3]
y_val_expected = [3, 2, 1, 0, 3, 2, 1, 0]