Searched defs:cluster_func (Results 1 – 8 of 8) sorted by relevance
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/ |
D | tpu_resource_read_for_write.cc | 49 tf_device::ClusterFuncOp cluster_func, Value result) { in GetResourceWriteResult() 72 bool ClusterFuncHasResourceRead(tf_device::ClusterFuncOp cluster_func, in ClusterFuncHasResourceRead() 84 getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation() 91 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) { in runOnOperation() local
|
D | tpu_resource_partitioning.cc | 74 void PartitionResourceReadsWrites(tf_device::ClusterFuncOp cluster_func) { in PartitionResourceReadsWrites() 144 getFunction()->walk([&cluster_funcs](tf_device::ClusterFuncOp cluster_func) { in runOnFunction() 147 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) in runOnFunction() local
|
D | tpu_dynamic_padding_mapper.cc | 59 tf_device::ClusterFuncOp cluster_func, tf_device::ReplicateOp replicate, in GetRemappedReplicatedInputIndices() 85 tf_device::ClusterFuncOp cluster_func, in GetRemappedPaddings() 181 LogicalResult RemapAndAssignPaddingMaps(tf_device::ClusterFuncOp cluster_func, in RemapAndAssignPaddingMaps() 211 module.walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
|
D | tpu_sharding_identification_pass.cc | 121 tf_device::ClusterFuncOp cluster_func, FuncOp func, Builder* builder) { in IdentifyXlaShardingForComputationInputs() 248 tf_device::ClusterFuncOp cluster_func, FuncOp func, Builder* builder) { in IdentifyXlaShardingForComputationOutputs() 305 Builder* builder, tf_device::ClusterFuncOp cluster_func) { in IdentifyXlaShardingForTPUComputation() 330 getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
|
D | annotate_parameter_replication.cc | 60 m.walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
|
D | tpu_space_to_depth_pass.cc | 423 TF::SpaceToDepthOp BuildSpaceToDepth(tf_device::ClusterFuncOp cluster_func, in BuildSpaceToDepth() 440 tf_device::ClusterFuncOp cluster_func, in HandleHostInput() 452 tf_device::ClusterFuncOp cluster_func, in HandleHostReplicatedInputs() 480 void HandleCluster(tf_device::ClusterFuncOp cluster_func, int32_t block_size, in HandleCluster() 650 Optional<tf_device::ClusterFuncOp> cluster_func; in runOnOperation() local
|
D | tpu_rewrite_pass.cc | 350 tf_device::ClusterFuncOp cluster_func, int num_replicas, in BuildCompileOp() 451 llvm::ArrayRef<Value> inputs, tf_device::ClusterFuncOp cluster_func, in BuildExecuteOp() 472 Operation* compile_op, tf_device::ClusterFuncOp cluster_func, in BuildParallelExecuteOp() 567 tf_device::ClusterFuncOp cluster_func, in Rewrite()
|
/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/ |
D | xla_sharding_util.cc | 212 mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder, in ExtractInputsForLogicalDevices() 313 mlir::tf_device::ClusterFuncOp cluster_func, in ParseAndValidateOutputSharding() 502 mlir::tf_device::ClusterFuncOp cluster_func, in GetOutputTypesForLogicalDeviceComputation() 539 mlir::tf_device::ClusterFuncOp cluster_func, in RemapOutputsFromLogicalDevices()
|