(self, arg0, arg1, global_jit_level)
| 482 | |
| 483 | # Runs a simple test with the input jit_level and fusion_only flag. |
| 484 | def simpleTest(self, arg0, arg1, global_jit_level): |
| 485 | config = config_pb2.ConfigProto() |
| 486 | config.graph_options.optimizer_options.global_jit_level = global_jit_level |
| 487 | |
| 488 | with session_lib.Session(config=config) as sess: |
| 489 | a1 = array_ops.placeholder(dtypes.float32, [2, 2], name="a1") |
| 490 | a2 = array_ops.placeholder(dtypes.float32, [2, 2], name="a2") |
| 491 | # Two element-wise ops. We need at least two ops since single |
| 492 | # element clusters are not passed to XLA in fusion_only mode. |
| 493 | a3 = a1 * a2 |
| 494 | a4 = a3 + a1 |
| 495 | # A matmul to break XLA clustering. |
| 496 | a5 = math_ops.matmul(a4, a1) |
| 497 | # Two more element-wise ops. |
| 498 | a6 = a5 - a4 |
| 499 | a7 = a6 + a2 |
| 500 | |
| 501 | run_metadata = config_pb2.RunMetadata() |
| 502 | output = test_utils.RunWithWarmup( |
| 503 | sess, |
| 504 | a7, { |
| 505 | a1: arg0, |
| 506 | a2: arg1 |
| 507 | }, |
| 508 | run_metadata=run_metadata, |
| 509 | options=config_pb2.RunOptions( |
| 510 | trace_level=config_pb2.RunOptions.FULL_TRACE)) |
| 511 | |
| 512 | labels = RunMetadataLabels(run_metadata) |
| 513 | |
| 514 | xla_compile_count = sum("XlaCompile(" in x for x in labels) |
| 515 | xla_run_count = sum("XlaRun(" in x for x in labels) |
| 516 | self.assertEqual(xla_compile_count, xla_run_count) |
| 517 | |
| 518 | return output, xla_run_count |
| 519 | |
| 520 | |
| 521 | class LazyCompilationTest(test.TestCase): |
nothing calls this directly
no test coverage detected