Home
last modified time | relevance | path

Searched refs:cluster_func (Results 1 – 19 of 19) sorted by relevance

/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
Dtpu_resource_read_for_write.cc49 tf_device::ClusterFuncOp cluster_func, Value result) { in GetResourceWriteResult() argument
62 if (assign_var_user.value().getDefiningOp() == cluster_func) in GetResourceWriteResult()
72 bool ClusterFuncHasResourceRead(tf_device::ClusterFuncOp cluster_func, in ClusterFuncHasResourceRead() argument
77 if (read_user == cluster_func) return true; in ClusterFuncHasResourceRead()
84 getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation() argument
85 cluster_funcs.push_back(cluster_func); in runOnOperation()
91 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) { in runOnOperation() local
92 builder.setInsertionPoint(cluster_func); in runOnOperation()
95 for (Value result : cluster_func.getResults()) { in runOnOperation()
97 auto resource_and_type = GetResourceWriteResult(cluster_func, result); in runOnOperation()
[all …]
Dtpu_dynamic_padding_mapper.cc59 tf_device::ClusterFuncOp cluster_func, tf_device::ReplicateOp replicate, in GetRemappedReplicatedInputIndices() argument
64 for (auto operand_and_idx : llvm::enumerate(cluster_func.getOperands())) { in GetRemappedReplicatedInputIndices()
85 tf_device::ClusterFuncOp cluster_func, in GetRemappedPaddings() argument
97 Attribute padding_map_attr = cluster_func->getAttr(kPaddingMapAttr); in GetRemappedPaddings()
102 return cluster_func.emitOpError() in GetRemappedPaddings()
110 return cluster_func.emitOpError( in GetRemappedPaddings()
116 return cluster_func.emitOpError(llvm::formatv( in GetRemappedPaddings()
123 return cluster_func.emitOpError() in GetRemappedPaddings()
128 return cluster_func.emitOpError() in GetRemappedPaddings()
137 cluster_func.emitWarning(llvm::formatv( in GetRemappedPaddings()
[all …]
Dtpu_rewrite_pass.cc350 tf_device::ClusterFuncOp cluster_func, int num_replicas, in BuildCompileOp() argument
357 cluster_func, num_replicas, num_cores_per_replica, in BuildCompileOp()
372 compile_op_operands.reserve(cluster_func.getNumOperands()); in BuildCompileOp()
374 for (auto operand_and_idx : llvm::enumerate(cluster_func.getOperands())) { in BuildCompileOp()
381 cluster_func.getLoc(), in BuildCompileOp()
387 FlatSymbolRefAttr func_attr = cluster_func.funcAttr(); in BuildCompileOp()
388 FuncOp func = cluster_func->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>( in BuildCompileOp()
400 cluster_func.getLoc(), in BuildCompileOp()
451 llvm::ArrayRef<Value> inputs, tf_device::ClusterFuncOp cluster_func, in BuildExecuteOp() argument
457 core_id, output_sharding_config, cluster_func, &output_types); in BuildExecuteOp()
[all …]
Dtpu_resource_partitioning.cc74 void PartitionResourceReadsWrites(tf_device::ClusterFuncOp cluster_func) { in PartitionResourceReadsWrites() argument
76 if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>( in PartitionResourceReadsWrites()
82 OpBuilder builder(cluster_func); in PartitionResourceReadsWrites()
87 for (Value result : cluster_func.results()) { in PartitionResourceReadsWrites()
104 cluster_func->getLoc(), partitioned_output_types, result, in PartitionResourceReadsWrites()
115 for (OpOperand& operand : cluster_func->getOpOperands()) { in PartitionResourceReadsWrites()
144 getFunction()->walk([&cluster_funcs](tf_device::ClusterFuncOp cluster_func) { in runOnFunction() argument
145 cluster_funcs.push_back(cluster_func); in runOnFunction()
147 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) in runOnFunction() local
148 PartitionResourceReadsWrites(cluster_func); in runOnFunction()
Dtpu_sharding_identification_pass.cc121 tf_device::ClusterFuncOp cluster_func, FuncOp func, Builder* builder) { in IdentifyXlaShardingForComputationInputs() argument
146 llvm::zip(cluster_func.operands(), function_block.getArguments())) { in IdentifyXlaShardingForComputationInputs()
182 cluster_func->setAttr(tensorflow::kInputShardingAttr, in IdentifyXlaShardingForComputationInputs()
248 tf_device::ClusterFuncOp cluster_func, FuncOp func, Builder* builder) { in IdentifyXlaShardingForComputationOutputs() argument
264 llvm::zip(cluster_func.results(), terminator->getOpOperands())) { in IdentifyXlaShardingForComputationOutputs()
298 cluster_func->setAttr(tensorflow::kOutputShardingAttr, in IdentifyXlaShardingForComputationOutputs()
305 Builder* builder, tf_device::ClusterFuncOp cluster_func) { in IdentifyXlaShardingForTPUComputation() argument
307 FuncOp func = cluster_func->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>( in IdentifyXlaShardingForTPUComputation()
308 cluster_func.func()); in IdentifyXlaShardingForTPUComputation()
316 if (auto use_spmd_attr = cluster_func->getAttrOfType<BoolAttr>( in IdentifyXlaShardingForTPUComputation()
[all …]
Dannotate_parameter_replication.cc60 m.walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation() argument
61 auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>(); in runOnOperation()
72 auto func = llvm::cast<FuncOp>(m.lookupSymbol(cluster_func.func())); in runOnOperation()
73 for (auto entry : llvm::enumerate(cluster_func.getOperands())) { in runOnOperation()
Dtpu_space_to_depth_pass.cc423 TF::SpaceToDepthOp BuildSpaceToDepth(tf_device::ClusterFuncOp cluster_func, in BuildSpaceToDepth() argument
435 cluster_func.getLoc(), transform_result_type, input, block_size); in BuildSpaceToDepth()
440 tf_device::ClusterFuncOp cluster_func, in HandleHostInput() argument
444 BuildSpaceToDepth(cluster_func, input, block_size, input_shape); in HandleHostInput()
445 cluster_func.setOperand(index, space_to_depth); in HandleHostInput()
452 tf_device::ClusterFuncOp cluster_func, in HandleHostReplicatedInputs() argument
471 BuildSpaceToDepth(cluster_func, input, block_size, input_shape); in HandleHostReplicatedInputs()
480 void HandleCluster(tf_device::ClusterFuncOp cluster_func, int32_t block_size, in HandleCluster() argument
483 llvm::dyn_cast<tf_device::ReplicateOp>(cluster_func->getParentOp()); in HandleCluster()
486 for (auto input : llvm::enumerate(cluster_func.operands())) { in HandleCluster()
[all …]
Dtf_passes.td258 replaces the `tf_device.cluster` op with an equivalent `tf_device.cluster_func`
260 the newly created functions and associated `tf_device.cluster_func` ops.
278 %cluster = "tf_device.cluster_func"(%arg0) {func = @_func} : (tensor<i32>) -> tensor<i32>
307 %0 = "tf_device.cluster_func"(%value) {func = @cluster} : (tensor<i32>) -> tensor<i32>
323 …%0 = "tf_device.cluster_func"(%value, %resource_read) {func = @cluster} : (tensor<i32>, tensor<i32…
563 …%computation = "tf_device.cluster_func"(%read) {func = @computation, use_spmd_for_xla_partitioning…
580 …%computation = "tf_device.cluster_func"(%partitioned_input) {func = @computation, use_spmd_for_xla…
597 let summary = "Rewrites a `tf_device.cluster_func` on TPUs into TPU runtime operations.";
600 This pass rewrites a `tf_device.cluster_func` operation into a sequence of `tf._TPUCompileMlir`
602 functionally equivalent to the function referenced by `tf_device.cluster_func`.
[all …]
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/
Dtpu_sharding_identification.mlir7 "tf_device.cluster_func"() {func = @empty_func, step_marker_location = ""} : () -> ()
24 …"tf_device.cluster_func"(%arg0) {func = @func_without_sharding, step_marker_location = ""} : (tens…
45 …"tf_device.cluster_func"(%arg0) {func = @func_without_sharding, step_marker_location = ""} : (tens…
66 …"tf_device.cluster_func"(%arg0) {func = @inputs_with_sharding_func, step_marker_location = ""} : (…
88 …"tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_sharding, step_marker_location = ""} : (…
113 …"tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_sharding_after_identity, step_marker_loc…
139 …"tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_sharding_after_read_variable, step_marke…
167 …"tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_sharding_after_cast, step_marker_locatio…
194 …"tf_device.cluster_func"(%arg0, %arg1) {func = @func_with_device_training_loop, step_marker_locati…
237 …"tf_device.cluster_func"(%arg0) {func = @cluster_func, step_marker_location = ""} : (tensor<*xi32>…
[all …]
Dtpu_dynamic_padding_mapper.mlir13 …"tf_device.cluster_func"(%ri_0, %ri_1) {func = @func0, padding_map = ["\10\02\18\01"]} : (tensor<i…
40 …"tf_device.cluster_func"(%ri_0, %ri_1, %ri_2) {func = @func1, padding_map = ["\10\02\18\01", "\10\…
72 …"tf_device.cluster_func"(%ri_0, %ri_1, %ri_2, %ri_3, %ri_4) {func = @func2, padding_map = ["\10\02…
93 …"tf_device.cluster_func"(%ri_1, %arg0, %ri_0) {func = @func3, padding_map = ["\10\02\18\01"]} : (t…
114 …"tf_device.cluster_func"(%arg0, %arg0, %arg0) {func = @func4, padding_map = ["\10\02\18\01"]} : (t…
128 …"tf_device.cluster_func"(%ri_1, %arg0, %ri_0) {func = @func5} : (tensor<i1>, tensor<i1>, tensor<i1…
144 …"tf_device.cluster_func"(%ri_1, %arg0, %ri_0) {func = @func6, padding_map = []} : (tensor<i1>, ten…
165 …"tf_device.cluster_func"(%ri_1) {func = @func7, padding_map = ["\10\02\18\01"]} : (tensor<i1>) -> …
192 …"tf_device.cluster_func"(%ri_0, %ri_2, %ri_3) {func = @func8, padding_map = ["\10\02\18\01", "\08\…
214 …"tf_device.cluster_func"(%ri_0, %ri_1) {func = @func9, padding_map = ["\10\02\18\01"]} : (tensor<i…
[all …]
Dtpu_resource_partitioning.mlir15 // CHECK: [[COMPUTATION:%.+]] = "tf_device.cluster_func"([[INPUT]])
16 …%2 = "tf_device.cluster_func"(%1) {func = @computation, use_spmd_for_xla_partitioning = true} : (t…
36 // CHECK: "tf_device.cluster_func"([[INPUT]])
37 …%2 = "tf_device.cluster_func"(%1) {func = @computation, use_spmd_for_xla_partitioning = true} : (t…
51 %2 = "tf_device.cluster_func"(%1) {func = @computation} : (tensor<i32>) -> tensor<i32>
55 …%5 = "tf_device.cluster_func"(%4) {func = @computation, use_spmd_for_xla_partitioning = false} : (…
63 %1 = "tf_device.cluster_func"(%0) {func = @computation} : (tensor<i32>) -> tensor<i32>
73 %1 = "tf_device.cluster_func"(%0) {func = @computation} : (tensor<i32>) -> tensor<i32>
85 %2 = "tf_device.cluster_func"(%1) {func = @computation} : (tensor<i32>) -> tensor<i32>
95 %2 = "tf_device.cluster_func"(%1) {func = @computation} : (tensor<i32>) -> tensor<i32>
[all …]
Dcluster_outlining.mlir13 …// CHECK: %[[CLUSTER_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[A_OUTPUT]]) {func = @[[CLUSTER:…
45 …// CHECK: %[[CLUSTER_0_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[A_OUTPUT]]) {func = @[[CLUSTE…
54 …// CHECK: %[[CLUSTER_1_OUTPUT:[0-9]*]] = "tf_device.cluster_func"(%[[CLUSTER_0_OUTPUT]], %[[D_OUTP…
89 …// CHECK: %[[CLUSTER_OUTPUT:[a-z0-9]*]], %{{.*}} = {{.*}} "tf_device.cluster_func"() {func = @[[CL…
107 // Tests cluster attributes are copied over to cluster_func.
118 // CHECK: "tf_device.cluster_func"
Dtpu_rewrite.mlir8 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
23 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
33 // Tests `tf_device.cluster_func` with missing `num_cores_per_replicas`
39 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, step_marker_location …
49 // Tests `tf_device.cluster_func` with bad `num_cores_per_replicas` attribute.
54 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
64 // Tests `tf_device.cluster_func` with missing `step_marker_location` attribute.
69 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
79 // Tests `tf_device.cluster_func` with bad `step_marker_location` attribute.
84 …"tf_device.cluster_func"() {_tpu_replicate = "cluster0", func = @empty_func, num_cores_per_replica…
[all …]
Dtpu-resource-read-for-write.mlir7 // CHECK-NEXT: [[CLUSTER:%.*]]:2 = "tf_device.cluster_func"([[ARG0]], [[ARG1]], [[READ]])
9 …%0:2 = "tf_device.cluster_func"(%arg0, %arg1) {_tpu_replicate = "write", func = @write_func} : (te…
26 …%1:2 = "tf_device.cluster_func"(%arg0, %arg1, %0) {_tpu_replicate = "read_write", func = @read_wri…
40 …%0:2 = "tf_device.cluster_func"(%arg0) {_tpu_replicate = "multiple_write", func = @multiple_write_…
55 …%0 = "tf_device.cluster_func"(%arg0) {_tpu_replicate = "multiple_uses", func = @multiple_result_us…
Dannotate-parameter-replication.mlir13 …%4 = "tf_device.cluster_func"(%ri_0, %3, %2) {func = @_func, device = ""} : (tensor<?xi32>, tensor…
49 …%2 = "tf_device.cluster_func"(%0, %1, %ri_2) {func = @_func, device = ""} : (tensor<?xi32>, tensor…
75 …%2 = "tf_device.cluster_func"(%0, %1) {func = @_func, device = ""} : (tensor<?xi32>, tensor<?xi32>…
Dtpu_space_to_depth_pass.mlir24 …%7:2 = "tf_device.cluster_func"(%1, %3, %5, %6) {_tpu_replicate = "while/cluster_while_body_271", …
106 …%9:4 = "tf_device.cluster_func"(%arg13, %arg14, %2, %4, %3, %5, %6, %7, %8) {_tpu_replicate = "clu…
180 …%9:4 = "tf_device.cluster_func"(%arg13, %arg14, %2, %4, %3, %5, %6, %7, %8) {_tpu_replicate = "clu…
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
Dxla_sharding_util.cc212 mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder, in ExtractInputsForLogicalDevices() argument
220 cluster_func.getOperands()); in ExtractInputsForLogicalDevices()
222 cluster_func.getOperation()->getAttrOfType<mlir::ArrayAttr>( in ExtractInputsForLogicalDevices()
247 return cluster_func.emitError( in ExtractInputsForLogicalDevices()
261 return cluster_func->emitOpError() in ExtractInputsForLogicalDevices()
289 cluster_func.getLoc(), sharding, input_value, builder, &tiled_inputs); in ExtractInputsForLogicalDevices()
313 mlir::tf_device::ClusterFuncOp cluster_func, in ParseAndValidateOutputSharding() argument
315 output_sharding_list->reserve(cluster_func.getNumResults()); in ParseAndValidateOutputSharding()
318 cluster_func.getOperation()->getAttrOfType<mlir::ArrayAttr>( in ParseAndValidateOutputSharding()
321 return cluster_func.emitError( in ParseAndValidateOutputSharding()
[all …]
Dxla_sharding_util.h41 mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder,
48 mlir::tf_device::ClusterFuncOp cluster_func,
56 mlir::tf_device::ClusterFuncOp cluster_func,
64 mlir::tf_device::ClusterFuncOp cluster_func,
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/ir/
Dtf_device_ops.td334 def TfDevice_ClusterFuncOp : TfDevice_Op<"cluster_func", []> {
336 The `tf_device.cluster_func` launches a function containing the body of a