/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
D | tpu_parallel_execute_sink_resource_write.cc | 49 tf_device::ParallelExecuteOp parallel_execute, Value result) { in GetSingleUseResourceWrite() argument 59 if (resource_handle_op == parallel_execute) return nullptr; in GetSingleUseResourceWrite() 63 parallel_execute.getOperation()->getBlock() && in GetSingleUseResourceWrite() 64 parallel_execute.getOperation()->isBeforeInBlock(resource_handle_op)) in GetSingleUseResourceWrite() 74 tf_device::ParallelExecuteOp parallel_execute) { in SinkResourceWritesIntoParallelExecute() argument 76 const int num_regions = parallel_execute.getNumRegions(); in SinkResourceWritesIntoParallelExecute() 85 Block& block = parallel_execute.GetRegionBlockWithIndex(i); in SinkResourceWritesIntoParallelExecute() 86 auto results = parallel_execute.GetRegionOutputs(i); in SinkResourceWritesIntoParallelExecute() 92 GetSingleUseResourceWrite(parallel_execute, result.value()); in SinkResourceWritesIntoParallelExecute() 115 parallel_execute.GetRegionBlockWithIndex(results_to_remove.index()); in SinkResourceWritesIntoParallelExecute() [all …]
|
D | tpu_merge_variables_with_execute.cc | 140 auto parallel_execute = llvm::dyn_cast<tf_device::ParallelExecuteOp>( in BuildVariableAccessInfo() local 143 parallel_execute ? parallel_execute.getOperation() : execute_launch; in BuildVariableAccessInfo() 241 parallel_execute in BuildVariableAccessInfo() 242 ? parallel_execute.GetRegionOutputs( in BuildVariableAccessInfo() 317 tf_device::ParallelExecuteOp parallel_execute, int start, in AppendTypes() argument 321 Block& block = parallel_execute.GetRegionBlockWithIndex(index); in AppendTypes() 331 void ReplaceParallelExecute(tf_device::ParallelExecuteOp parallel_execute, in ReplaceParallelExecute() argument 336 Operation* parallel_execute_op = parallel_execute.getOperation(); in ReplaceParallelExecute() 346 AppendTypes(&output_types, parallel_execute, 0, region_index); in ReplaceParallelExecute() 352 &output_types, parallel_execute, region_index + 1, num_regions); in ReplaceParallelExecute() [all …]
|
D | tf_passes.td | 339 … = "Extracts TPU outside compilation computation to a separate tf_device.parallel_execute region."; 345 a tf_device.parallel_execute region. The TPU cluster is also moved to a 346 tf_device.parallel_execute region. Communication ops between device and host are 363 will become a tf_device.parallel_execute op with a CPU/host region and 368 %0 = "tf_device.parallel_execute"() ( { 693 %1 = "tf_device.parallel_execute"() ( {
|
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/ |
D | tpu_parallel_execute_sink_resource_write.mlir | 6 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute" 7 %0:2 = "tf_device.parallel_execute"() ( { 22 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute" 23 %0:2 = "tf_device.parallel_execute"() ( { 37 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute" 38 %0:2 = "tf_device.parallel_execute"() ( { 51 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute" 52 %0:2 = "tf_device.parallel_execute"() ( { 64 // CHECK: [[PARALLEL_EXECUTE:%.+]]:2 = "tf_device.parallel_execute" 65 %0:2 = "tf_device.parallel_execute"() ( { [all …]
|
D | tf_device_ops_invalid.mlir | 200 // Check that a parallel_execute op with a single region is not allowed. 202 "tf_device.parallel_execute"() ( { 203 // expected-error@-1 {{'tf_device.parallel_execute' op must have at least two regions.}} 211 // Check that a parallel_execute op with empty regions are not allowed. 213 "tf_device.parallel_execute"() ( { 214 // expected-error@-1 {{'tf_device.parallel_execute' op region #0 ('regions') failed to verify const… 224 // Check that a parallel_execute ops with invalid number of output types are 227 "tf_device.parallel_execute"() ({ 228 // expected-error@-1 {{'tf_device.parallel_execute' op number of output types (3) must match the to… 242 // Check that a parallel_execute ops with mismatching output types are not [all …]
|
D | parallel_execute_to_islands.mlir | 7 "tf_device.parallel_execute"() ({ 35 %4:2 = "tf_device.parallel_execute"() ({ 68 %3:2 = "tf_device.parallel_execute"() ({ 96 %1:2 = "tf_device.parallel_execute"() ({ 128 %2:2 = "tf_device.parallel_execute"() ({
|
D | tpu_extract_outside_compilation.mlir | 6 // Tests that TPU cluster with no outside compilation does not generate parallel_execute. 18 // CHECK-NOT: "tf_device.parallel_execute" 24 // CHECK: "tf_device.parallel_execute" 46 // CHECK: "tf_device.parallel_execute" 71 // CHECK: "tf_device.parallel_execute" 94 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" 122 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:2 = "tf_device.parallel_execute" 147 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" 179 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" 201 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" [all …]
|
D | tpu-merge-variables-with-execute.mlir | 229 // with TPUExecutes in a tf_device.parallel_execute. 231 // CHECK-LABEL: func @parallel_execute 235 func @parallel_execute( 242 // CHECK: "tf_device.parallel_execute" 243 %pe:2 = "tf_device.parallel_execute"() ( { 273 // tf_device.parallel_execute that is replicated (tf_device.replicate). 295 // CHECK: "tf_device.parallel_execute" 296 %pe:2 = "tf_device.parallel_execute"() ( {
|
D | tf_device_ops.mlir | 167 "tf_device.parallel_execute"() ({ 178 "tf_device.parallel_execute"() ({ 192 "tf_device.parallel_execute"() ({
|
D | tpu-dynamic-layout-pass.mlir | 410 // CHECK-LABEL: func @parallel_execute 411 func @parallel_execute(%arg0: tensor<*x!tf.resource> {tf.device = "/device:CPU:0"}) { 427 // CHECK: "tf_device.parallel_execute" 428 "tf_device.parallel_execute"() ({ 502 // CHECK: "tf_device.parallel_execute" 503 "tf_device.parallel_execute"() ({
|
D | tpu_rewrite.mlir | 1229 // CHECK: [[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:2 = "tf_device.parallel_execute" 1267 // CHECK: [[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:2 = "tf_device.parallel_execute" 1391 // parallel_execute. 1403 // CHECK: "tf_device.parallel_execute" 1409 %3 = "tf_device.parallel_execute"() ( { 1448 // CHECK: "tf_device.parallel_execute" 1512 // CHECK: "tf_device.parallel_execute" 1548 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" 1586 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]] = "tf_device.parallel_execute" 1623 // CHECK: %[[PARALLEL_EXECUTE_OUTPUT:[0-9]*]]:3 = "tf_device.parallel_execute" [all …]
|
D | tpu-variable-runtime-reformatting.mlir | 204 "tf_device.parallel_execute"() ({
|
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
D | xla_sharding_util.cc | 393 mlir::tf_device::ParallelExecuteOp parallel_execute) { in GetTileShardedOutputsToMerge() argument 401 const auto output_from_logical_device = parallel_execute.GetRegionOutputs( in GetTileShardedOutputsToMerge() 413 mlir::tf_device::ParallelExecuteOp parallel_execute, in HandleTileShardedOutputs() argument 417 builder->setInsertionPointAfter(parallel_execute); in HandleTileShardedOutputs() 422 cluster_func_output_index, sharding, parallel_execute); in HandleTileShardedOutputs() 540 mlir::tf_device::ParallelExecuteOp parallel_execute, in RemapOutputsFromLogicalDevices() argument 576 parallel_execute.GetRegionOutputs( in RemapOutputsFromLogicalDevices() 584 output_index, output_sharding, parallel_execute); in RemapOutputsFromLogicalDevices() 594 cluster_func_output, parallel_execute, in RemapOutputsFromLogicalDevices() 608 const auto output_from_logical_device = parallel_execute.GetRegionOutputs( in RemapOutputsFromLogicalDevices()
|
D | xla_sharding_util.h | 65 mlir::tf_device::ParallelExecuteOp parallel_execute,
|
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/ |
D | tf_device_ops.td | 130 def TfDevice_ParallelExecuteOp : TfDevice_Op<"parallel_execute", 165 // Checks if a tf_device.parallel_execute index'th region block wraps a 190 A tf_device.parallel_execute inside the tf_device.replicate op region may be 193 concurrent execution (i.e. region) defined by tf_device.parallel_execute op.
|