diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD index 5251a2953e7092..00cb78d3a71774 100644 --- a/tensorflow/compiler/tests/BUILD +++ b/tensorflow/compiler/tests/BUILD @@ -594,6 +594,7 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/framework:errors", "//tensorflow/python/framework:tensor_shape", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:linalg_ops", @@ -1371,6 +1372,7 @@ tf_xla_py_strict_test( ], deps = [ ":xla_test", + "//tensorflow/python/framework:errors", "//tensorflow/python/framework:test_lib", "//tensorflow/python/ops:array_ops", "//tensorflow/python/ops:linalg_ops", diff --git a/tensorflow/compiler/tests/qr_op_test.py b/tensorflow/compiler/tests/qr_op_test.py index 4eac4a970bf199..b77bc9c46aac28 100644 --- a/tensorflow/compiler/tests/qr_op_test.py +++ b/tensorflow/compiler/tests/qr_op_test.py @@ -21,6 +21,7 @@ import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.python.framework import errors from tensorflow.python.framework import test_util from tensorflow.python.ops import array_ops from tensorflow.python.ops import linalg_ops @@ -143,6 +144,18 @@ def testRepeatedColumn(self, rows, cols): x_np[:, 1] = x_np[:, 2] self._test(x_np, full_matrices=True, full_rank=False) + def testVectorInputRaisesError(self): + # Regression test for GitHub issue 110798. The graph-level shape check + # only runs for inputs of known rank, so a rank-1 input reaching the + # compiler through an unknown-rank placeholder must be rejected at + # compile time instead of reaching a fatal check inside QrExplicit. + with self.session() as sess: + x_tf = array_ops.placeholder(np.float32) + with self.device_scope(): + q_tf, r_tf = linalg_ops.qr(x_tf, full_matrices=True) + with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"): + sess.run([q_tf, r_tf], feed_dict={x_tf: np.zeros([8], np.float32)}) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tests/svd_op_test.py b/tensorflow/compiler/tests/svd_op_test.py index 356a4bc7715b0a..7214231579b439 100644 --- a/tensorflow/compiler/tests/svd_op_test.py +++ b/tensorflow/compiler/tests/svd_op_test.py @@ -20,6 +20,7 @@ import numpy as np from tensorflow.compiler.tests import xla_test +from tensorflow.python.framework import errors from tensorflow.python.framework import tensor_shape from tensorflow.python.ops import array_ops from tensorflow.python.ops import gen_linalg_ops @@ -86,6 +87,18 @@ def testSvd(self, n, dtype): self._testSvdCorrectness(dtype, batch_dims + (2 * n, n)) self._testSvdCorrectness(dtype, batch_dims + (n, 2 * n)) + def testVectorInputRaisesError(self): + # Like GitHub issue 110798 for QR: the graph-level shape check only runs + # for inputs of known rank, so a rank-1 input reaching the compiler + # through an unknown-rank placeholder must be rejected at compile time + # instead of reading a negative trailing dimension. + with self.session() as sess: + x_tf = array_ops.placeholder(np.float32) + with self.test_scope(): + s, u, v = linalg_ops.svd(x_tf, full_matrices=True) + with self.assertRaisesRegex(errors.InvalidArgumentError, "rank >= 2"): + sess.run([s, u, v], feed_dict={x_tf: np.zeros([8], np.float32)}) + if __name__ == "__main__": test.main() diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD index df541ca511272c..f1b6ab5f79ecb4 100644 --- a/tensorflow/compiler/tf2xla/kernels/BUILD +++ b/tensorflow/compiler/tf2xla/kernels/BUILD @@ -850,6 +850,8 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", "//tensorflow/core:lib", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", "@xla//xla:shape_util", "@xla//xla:xla_data_proto_cc", "@xla//xla/hlo/builder/lib:constants", @@ -1563,6 +1565,8 @@ tf_kernel_library( "//tensorflow/compiler/tf2xla:xla_resource", "//tensorflow/compiler/tf2xla/ops:xla_ops", "//tensorflow/core:framework", + "@com_google_absl//absl/status", + "@com_google_absl//absl/strings", "@xla//xla/hlo/builder:xla_builder", "@xla//xla/hlo/builder/lib:qr", ], diff --git a/tensorflow/compiler/tf2xla/kernels/qr_op.cc b/tensorflow/compiler/tf2xla/kernels/qr_op.cc index 6120903fe9c991..4f8646be8683dc 100644 --- a/tensorflow/compiler/tf2xla/kernels/qr_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/qr_op.cc @@ -13,12 +13,15 @@ See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/qr.h" #include "xla/hlo/builder/xla_builder.h" #include "tensorflow/core/framework/op_kernel.h" #include "tensorflow/core/framework/op_requires.h" +#include "tensorflow/core/framework/tensor_shape.h" namespace tensorflow { namespace { @@ -29,6 +32,14 @@ class QROp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_)); } void Compile(XlaOpKernelContext* ctx) override { + // The rank is only known at compile time when the graph-level shape + // inference saw an unknown rank, so it must be validated here before + // QrExplicit reads the trailing two dimensions. + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); xla::XlaOp q, r; xla::QrExplicit(ctx->Input(0), full_matrices_, q, r); ctx->SetOutput(0, q); diff --git a/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc b/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc index 9583341992ce7d..941b6785dc3150 100644 --- a/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc +++ b/tensorflow/compiler/tf2xla/kernels/xla_svd_op.cc @@ -17,6 +17,8 @@ limitations under the License. #include #include +#include "absl/status/status.h" +#include "absl/strings/str_cat.h" #include "tensorflow/compiler/tf2xla/xla_op_kernel.h" #include "tensorflow/compiler/tf2xla/xla_op_registry.h" #include "xla/hlo/builder/lib/constants.h" @@ -51,6 +53,11 @@ class XlaSvdOp : public XlaOpKernel { } } void Compile(XlaOpKernelContext* ctx) override { + const TensorShape input_shape = ctx->InputShape(0); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); auto result = xla::SVD(ctx->Input(0), max_iter_, epsilon_, precision_config_.operand_precision(0)); ctx->SetOutput(0, result.d); @@ -71,7 +78,14 @@ class SvdOp : public XlaOpKernel { OP_REQUIRES_OK(ctx, ctx->GetAttr("full_matrices", &full_matrices_)); } void Compile(XlaOpKernelContext* ctx) override { + // The rank is only known at compile time when the graph-level shape + // inference saw an unknown rank, so it must be validated here before + // the trailing two dimensions are read. const TensorShape input_shape = ctx->InputShape("input"); + OP_REQUIRES(ctx, input_shape.dims() >= 2, + absl::InvalidArgumentError( + absl::StrCat("Input must have rank >= 2, got shape ", + input_shape.DebugString()))); int m = input_shape.dim_size(input_shape.dims() - 2); int n = input_shape.dim_size(input_shape.dims() - 1); // This is based on heuristics that approx log(n) sweep updates are needed. diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc index 94bdac9439a21b..0b1f2b783b05bd 100644 --- a/tensorflow/lite/kernels/transpose.cc +++ b/tensorflow/lite/kernels/transpose.cc @@ -50,26 +50,12 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context, // Ensure validity of the permutations tensor as a 1D tensor. TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1); TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims); - // `perm` must be a permutation of [0, dims), not merely a set of in-range - // values: the output shape and the element offsets are both derived from it, - // and a repeated entry makes them disagree with the input extent. - // `dims` is bounded by kTransposeMaxDimensions, which Prepare() enforces - // before this function is reachable, so a 64-bit mask is sufficient and - // avoids a heap allocation in the kernel. - static_assert(kTransposeMaxDimensions <= 64, - "Permutation bitmask assumes at most 64 dimensions."); - uint64_t seen = 0; for (int idx = 0; idx < dims; ++idx) { TF_LITE_ENSURE_MSG(context, (perm_data[idx] >= -dims && perm_data[idx] < dims), "Transpose op permutations array is out of bounds."); new_perm_data[idx] = perm_data[idx]; if (new_perm_data[idx] < 0) new_perm_data[idx] += dims; - const uint64_t bit = uint64_t{1} << new_perm_data[idx]; - TF_LITE_ENSURE_MSG( - context, (seen & bit) == 0, - "Transpose op permutations array must not contain duplicate values."); - seen |= bit; } // Determine size of output tensor. diff --git a/tensorflow/lite/kernels/transpose_test.cc b/tensorflow/lite/kernels/transpose_test.cc index 11930f342ce91a..601e8bf4355540 100644 --- a/tensorflow/lite/kernels/transpose_test.cc +++ b/tensorflow/lite/kernels/transpose_test.cc @@ -193,23 +193,6 @@ TEST(TransposeTest, TestPermOutOfBounds) { EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}), "Transpose op permutations array is out of bounds."); } - -// `perm` must be a permutation of [0, dims). A repeated entry passes the range -// check but makes the derived output shape disagree with the input extent, so -// the kernel reads outside the input tensor. -TEST(TransposeTest, TestPermDuplicateValues) { - EXPECT_DEATH( - TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 2}), - "Transpose op permutations array must not contain duplicate values."); -} - -// Duplicates must also be rejected after negative entries are normalised: -// on a rank-2 input {0, -2} normalises to {0, 0}. -TEST(TransposeTest, TestPermDuplicateValuesAfterNegativeNormalization) { - EXPECT_DEATH( - TransposeOpConstModel({2, 3}, {2}, {0, -2}), - "Transpose op permutations array must not contain duplicate values."); -} #endif TEST(TransposeTest, TestInt41DInputConstTensor) { diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py index d7ed4722770933..cdad24086e2d69 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py @@ -529,25 +529,41 @@ def f(a, b): @tf_export.tf_export('experimental.numpy.logaddexp', v1=[]) @np_utils.np_doc('logaddexp') def logaddexp(x1, x2): - amax = maximum(x1, x2) - delta = x1 - x2 - return np_array_ops.where( - isnan(delta), - x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp(-abs(delta))), - ) + + def f(x1, x2): + if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact): + float_dtype = np_utils.result_type(float) + x1 = math_ops.cast(x1, float_dtype) + x2 = math_ops.cast(x2, float_dtype) + amax = maximum(x1, x2) + delta = x1 - x2 + return np_array_ops.where( + isnan(delta), + x1 + x2, # NaNs or infinities of the same sign. + amax + log1p(exp(-abs(delta))), + ) + + return _bin_op(f, x1, x2) @tf_export.tf_export('experimental.numpy.logaddexp2', v1=[]) @np_utils.np_doc('logaddexp2') def logaddexp2(x1, x2): - amax = maximum(x1, x2) - delta = x1 - x2 - return np_array_ops.where( - isnan(delta), - x1 + x2, # NaNs or infinities of the same sign. - amax + log1p(exp2(-abs(delta))) / np.log(2), - ) + + def f(x1, x2): + if not np.issubdtype(x1.dtype.as_numpy_dtype, np.inexact): + float_dtype = np_utils.result_type(float) + x1 = math_ops.cast(x1, float_dtype) + x2 = math_ops.cast(x2, float_dtype) + amax = maximum(x1, x2) + delta = x1 - x2 + return np_array_ops.where( + isnan(delta), + x1 + x2, # NaNs or infinities of the same sign. + amax + log1p(exp2(-abs(delta))) / math_ops.cast(np.log(2), x1.dtype), + ) + + return _bin_op(f, x1, x2) @tf_export.tf_export('experimental.numpy.polyval', v1=[]) diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py index 18c50d31fbdbce..d963f676d09cae 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_no_numpy_methods_test.py @@ -126,6 +126,28 @@ def testAverageAcceptsIntegerInputs(self): err_msg='average({})'.format(arg), ) + def testLogaddexpAcceptsIntegerInputs(self): + # `logaddexp` and `logaddexp2` promote integer arguments to a float dtype, + # which must not depend on `enable_numpy_methods_on_tensor()`. + x1 = np.array([1, 2, 3], dtype=np.int32) + for x2 in (np.int32(1), np.array([4, 5, 6], dtype=np.int32)): + actual = np_math_ops.logaddexp(x1, x2) + np.testing.assert_allclose( + np.asarray(actual), + np.logaddexp(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='logaddexp({}, {})'.format(x1, x2), + ) + actual2 = np_math_ops.logaddexp2(x1, x2) + np.testing.assert_allclose( + np.asarray(actual2), + np.logaddexp2(x1, x2), + rtol=1e-6, + atol=1e-6, + err_msg='logaddexp2({}, {})'.format(x1, x2), + ) + def testFloatInputsAreUnchanged(self): arg = np.array([1.5, 2.5, 3.5], dtype=np.float32) for name, tf_fun, np_fun in _PROMOTING_UNARY_OPS: diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py index 26f33fd00c3b19..54b71aa89f8c66 100644 --- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py +++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py @@ -207,6 +207,30 @@ def testHypotSpecialCases(self): expected = np.hypot(x, y) np.testing.assert_equal(actual.tolist(), expected.tolist()) + def testLogaddexp(self): + self._testBinaryOp(np_math_ops.logaddexp, np.logaddexp, 'logaddexp') + self._testBinaryOp(np_math_ops.logaddexp2, np.logaddexp2, 'logaddexp2') + + def testLogaddexpNonFloatInputs(self): + int_args = [ + ([1, 2, 3], [4, 5, 6]), + (np.array([1, 2], dtype=np.int32), np.array([3, 4], dtype=np.int32)), + (np.array([1, 2], dtype=np.int64), np.array([3, 4], dtype=np.int64)), + (1, 2), + ([1.0, 2.0], [3.0, 4.0]), + ] + for x1, x2 in int_args: + self.match( + np_math_ops.logaddexp(x1, x2), + np.logaddexp(np.asarray(x1), np.asarray(x2)), + msg='logaddexp({}, {})'.format(x1, x2), + ) + self.match( + np_math_ops.logaddexp2(x1, x2), + np.logaddexp2(np.asarray(x1), np.asarray(x2)), + msg='logaddexp2({}, {})'.format(x1, x2), + ) + def match(self, actual, expected, msg='', check_dtype=True): self.assertIsInstance(actual, np_arrays.ndarray) if check_dtype: diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 9e402b8148fed0..39c41bbce90a62 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -2311,6 +2311,7 @@ xla_cc_test( srcs = ["hlo_schedule_test.cc"], deps = [ ":buffer_value", + ":hlo_module_config", "//xla:literal_util", "//xla:shape_util", "//xla:xla_data_proto_cc", @@ -2325,6 +2326,8 @@ xla_cc_test( "//xla/tsl/platform:statusor", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/log", + "@com_google_absl//absl/log:check", + "@com_google_benchmark//:benchmark", "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc index 74325b23ede671..947054a2c0d509 100644 --- a/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc +++ b/third_party/xla/xla/service/gpu/triton_fusion_analysis_test.cc @@ -683,27 +683,6 @@ e { ::testing::HasSubstr("Unsupported broadcast"))); } -TEST_F(TritonDotAnalysisTest, OutputBroadcastIsNotAccepted) { - ASSERT_OK_AND_ASSIGN(std::unique_ptr module, - ParseAndReturnVerifiedModule(R"( -HloModule t - -ENTRY e { - p0 = f16[2,35] parameter(0) - p0c = bf16[2,35] convert(p0) - p1 = bf16[35,2] parameter(1) - dot = bf16[2,2] dot(p0c, p1), - lhs_contracting_dims={1}, rhs_contracting_dims={0} - ROOT bc = bf16[2,2,100] broadcast(dot), dimensions={0,1} -})")); - EXPECT_TRUE(GemmFusion(se::GpuComputeCapability{se::CudaComputeCapability{ - se::CudaComputeCapability::kAmpere, 0}}) - .Run(module.get()) - .value()); - EXPECT_EQ(module->entry_computation()->root_instruction()->opcode(), - HloOpcode::kBroadcast); -} - TEST_F(TritonDotAnalysisTest, DegenerateSplitFragmentIsHandled) { ASSERT_OK_AND_ASSIGN(std::unique_ptr module, ParseAndReturnVerifiedModule(R"( diff --git a/third_party/xla/xla/service/hlo_schedule_test.cc b/third_party/xla/xla/service/hlo_schedule_test.cc index 212c7f6556cf62..c6567096e7eda9 100644 --- a/third_party/xla/xla/service/hlo_schedule_test.cc +++ b/third_party/xla/xla/service/hlo_schedule_test.cc @@ -15,6 +15,8 @@ limitations under the License. #include "xla/hlo/ir/hlo_schedule.h" +#include +#include #include #include #include @@ -22,7 +24,9 @@ limitations under the License. #include #include #include "absl/algorithm/container.h" +#include "absl/log/check.h" #include "absl/log/log.h" +#include "benchmark/benchmark.h" #include "xla/hlo/analysis/alias_info.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" @@ -33,6 +37,7 @@ limitations under the License. #include "xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h" #include "xla/literal_util.h" #include "xla/service/buffer_value.h" +#include "xla/service/hlo_module_config.h" #include "xla/shape.h" #include "xla/shape_util.h" #include "xla/tsl/lib/core/status_test_util.h" @@ -660,5 +665,156 @@ ENTRY %test (arg.0: (f32[], f32[])) -> f32[] { ASSERT_OK(module->schedule().Update()); ASSERT_OK(module->schedule().Verify()); } + +std::unique_ptr BuildBenchmarkModule(int64_t num_instructions) { + HloModuleConfig config; + auto module = std::make_unique("bm_module", config); + HloComputation::Builder builder("entry"); + Shape shape = ShapeUtil::MakeShape(F32, {}); + HloInstruction* c0 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + HloInstruction* c1 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(2.0f))); + std::vector sequence; + sequence.reserve(num_instructions + 100); + sequence.push_back(c0); + sequence.push_back(c1); + for (int64_t i = 2; i < num_instructions; ++i) { + HloInstruction* op0 = sequence[i - 1]; + HloInstruction* op1 = sequence[i / 2]; + HloInstruction* add = builder.AddInstruction( + HloInstruction::CreateBinary(shape, HloOpcode::kAdd, op0, op1)); + sequence.push_back(add); + } + HloComputation* entry = module->AddEntryComputation(builder.Build()); + CHECK_OK(module->set_schedule(HloSchedule(module.get()))); + module->schedule().set_sequence(entry, sequence); + return module; +} + +// Benchmarks HloSchedule::Update() when a small number of new instructions (5 +// copies) are inserted into an existing scheduled computation. +void BM_HloScheduleUpdate_InsertFewCopies(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + constexpr int64_t kNumCopies = 5; + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + HloComputation* entry = module->entry_computation(); + const auto& seq = module->schedule().sequence(entry).instructions(); + for (int k = 1; k <= kNumCopies; ++k) { + int64_t target_idx = (num_instructions * k) / (kNumCopies + 1); + HloInstruction* new_copy = + entry->AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kCopy, seq[target_idx - 1])); + CHECK_OK(seq[target_idx]->ReplaceOperandWith(0, new_copy)); + } + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_InsertFewCopies) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when a proportional number of instructions +// (~1%) are inserted into an existing scheduled computation. +void BM_HloScheduleUpdate_Insert1PercentCopies(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + const int64_t num_copies = std::max(1, num_instructions / 100); + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + HloComputation* entry = module->entry_computation(); + const auto& seq = module->schedule().sequence(entry).instructions(); + for (int k = 1; k <= num_copies; ++k) { + int64_t target_idx = (num_instructions * k) / (num_copies + 1); + HloInstruction* new_copy = + entry->AddInstruction(HloInstruction::CreateUnary( + shape, HloOpcode::kCopy, seq[target_idx - 1])); + CHECK_OK(seq[target_idx]->ReplaceOperandWith(0, new_copy)); + } + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_Insert1PercentCopies) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when instructions are removed from the +// computation. +void BM_HloScheduleUpdate_RemoveFewInstructions(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + constexpr int64_t kNumRemovals = 5; + HloModuleConfig config; + Shape shape = ShapeUtil::MakeShape(F32, {}); + + for (auto _ : state) { + state.PauseTiming(); + auto module = std::make_unique("bm_module", config); + HloComputation::Builder builder("entry"); + HloInstruction* c0 = builder.AddInstruction( + HloInstruction::CreateConstant(LiteralUtil::CreateR0(1.0f))); + std::vector sequence; + sequence.push_back(c0); + HloInstruction* prev = c0; + std::vector dead_instructions; + for (int64_t i = 1; i < num_instructions; ++i) { + if (dead_instructions.size() < kNumRemovals && + i % (num_instructions / (kNumRemovals + 1)) == 0) { + HloInstruction* dead = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, c0)); + sequence.push_back(dead); + dead_instructions.push_back(dead); + } else { + prev = builder.AddInstruction( + HloInstruction::CreateUnary(shape, HloOpcode::kNegate, prev)); + sequence.push_back(prev); + } + } + HloComputation* entry = module->AddEntryComputation(builder.Build()); + CHECK_OK(module->set_schedule(HloSchedule(module.get()))); + module->schedule().set_sequence(entry, sequence); + + for (HloInstruction* dead : dead_instructions) { + CHECK_OK(entry->RemoveInstruction(dead)); + } + + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_RemoveFewInstructions) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + +// Benchmarks HloSchedule::Update() when no modifications were made. +void BM_HloScheduleUpdate_NoModifications(benchmark::State& state) { + const int64_t num_instructions = state.range(0); + + for (auto _ : state) { + state.PauseTiming(); + auto module = BuildBenchmarkModule(num_instructions); + state.ResumeTiming(); + CHECK_OK(module->schedule().Update()); + } +} +BENCHMARK(BM_HloScheduleUpdate_NoModifications) + ->Arg(1000) + ->Arg(5000) + ->Arg(10000) + ->Arg(50000); + } // namespace } // namespace xla