STT-tensorflow/tensorflow/compiler/jit/xla_compilation_cache_test.cc
Jeffrey A. Dean ec030f72f3 Performance improvements to speed up invocation of XLA code, by making
canonicalization and signature generation faster

Added benchmark for XlaCompilationCache::BuildSignature to measure time
taken to build a signature for the cache.

Base is this CL with just the changes to add the benchmark in
xla_compilation_cache_test.cc, New is this whole CL.

Run on desktop machine (40 X 2793 MHz CPUs); 2019-09-17T08:30:04.125894664-07:00
CPU: Intel Ivybridge with HyperThreading (20 cores) dL1:32KB dL2:256KB dL3:25MB
Benchmark                                      Base (ns)    New (ns) Improvement
----------------------------------------------------------------------------
BM_BuildSignature/0                                  226          87    +61.5%
BM_BuildSignature/1                                  337         171    +49.3%
BM_BuildSignature/2                                  504         259    +48.6%
BM_BuildSignature/5                                 1008         592    +41.3%
BM_BuildSignature/10                                1751        1238    +29.3%

RELNOTES: n/a
PiperOrigin-RevId: 276289188
Change-Id: Ia47343203f6ac587a921a92f86c2428dd04db2a7
2019-10-23 09:39:27 -07:00

81 lines
3.0 KiB
C++

/* Copyright 2018 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/jit/xla_compilation_cache.h"
#include "tensorflow/compiler/tf2xla/shape_util.h"
#include "tensorflow/core/platform/test.h"
#include "tensorflow/core/platform/test_benchmark.h"
namespace tensorflow {
namespace {
TEST(XlaCompilationCacheTest, SignatureEquality) {
NameAttrList fn;
fn.set_name("afunction");
std::vector<XlaCompiler::Argument> args(1);
args[0].kind = XlaCompiler::Argument::kConstant;
args[0].type = DT_INT32;
args[0].shape = TensorShape({4, 0});
args[0].constant_value = Tensor(DT_INT32, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s1,
XlaCompilationCache::BuildSignature(fn, args));
args[0].type = DT_FLOAT;
args[0].constant_value = Tensor(DT_FLOAT, {4, 0});
TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s2,
XlaCompilationCache::BuildSignature(fn, args));
args[0].shape = TensorShape({0, 4});
args[0].constant_value = Tensor(DT_FLOAT, {0, 4});
TF_ASSERT_OK_AND_ASSIGN(XlaCompilationCache::Signature s3,
XlaCompilationCache::BuildSignature(fn, args));
std::vector<XlaCompilationCache::Signature> signatures = {s1, s2, s3};
for (int i = 0; i < signatures.size(); ++i) {
for (int j = 0; j < signatures.size(); ++j) {
EXPECT_EQ(i == j, signatures[i] == signatures[j])
<< signatures[i].HumanString() << " " << signatures[j].HumanString();
}
}
}
static void BM_BuildSignature(int iters, int n_args) {
NameAttrList fn;
fn.set_name("afunction");
for (int i = 0; i < n_args; i++) {
(*fn.mutable_attr())[absl::StrCat("T", i)].set_type(DT_FLOAT);
}
std::vector<XlaCompiler::Argument> args(n_args);
for (int i = 0; i < n_args; i++) {
args[i].kind = (((i % 3) == 0) ? XlaCompiler::Argument::kConstant
: XlaCompiler::Argument::kParameter);
args[i].type = DT_INT32;
args[i].shape = TensorShape({4, 0});
args[i].constant_value = Tensor(DT_INT32, {4, 0});
}
while (--iters > 0) {
xla::StatusOr<XlaCompilationCache::Signature> s =
XlaCompilationCache::BuildSignature(fn, args);
CHECK(s.ok());
XlaCompilationCache::Signature sig = std::move(s.ValueOrDie());
}
}
BENCHMARK(BM_BuildSignature)->Arg(0)->Arg(1)->Arg(2)->Arg(5)->Arg(10);
} // namespace
} // namespace tensorflow