Home
last modified time | relevance | path

Searched defs:cluster_func (Results 1 – 10 of 10) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtpu_resource_read_for_write.cc50 tf_device::ClusterFuncOp cluster_func, Value result) { in GetResourceWriteResult()
73 bool ClusterFuncHasResourceRead(tf_device::ClusterFuncOp cluster_func, in ClusterFuncHasResourceRead()
85 getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
92 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) { in runOnOperation() local
H A Dtpu_resource_partitioning.cc75 void PartitionResourceReadsWrites(tf_device::ClusterFuncOp cluster_func) { in PartitionResourceReadsWrites()
145 getOperation()->walk([&cluster_funcs](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
148 for (tf_device::ClusterFuncOp cluster_func : cluster_funcs) in runOnOperation() local
H A Dtpu_rewrite_pass.cc332 tf_device::ClusterFuncOp cluster_func, int num_replicas, in BuildCompileOp()
439 llvm::ArrayRef<Value> inputs, tf_device::ClusterFuncOp cluster_func, in BuildExecuteOp()
466 tf_device::ClusterFuncOp cluster_func, in MovePreservedParallelExecuteChildren()
510 Operation* compile_op, tf_device::ClusterFuncOp cluster_func, in AddToParallelExecuteOp()
614 tf_device::ClusterFuncOp cluster_func, OpBuilder* builder, in BuildSingletonParallelExecuteOp()
697 tf_device::ClusterFuncOp cluster_func, in Rewrite()
H A Dtpu_sharding_identification_pass.cc183 bool infer_from_computation, tf_device::ClusterFuncOp cluster_func, in IdentifyXlaShardingForComputationInputs()
353 bool infer_from_computation, tf_device::ClusterFuncOp cluster_func, in IdentifyXlaShardingForComputationOutputs()
413 Builder* builder, tf_device::ClusterFuncOp cluster_func) { in IdentifyXlaShardingForTPUComputation()
499 getOperation().walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
H A Dtpu_space_to_depth_pass.cc378 TF::SpaceToDepthOp BuildSpaceToDepth(tf_device::ClusterFuncOp cluster_func, in BuildSpaceToDepth()
395 tf_device::ClusterFuncOp cluster_func, in HandleHostInput()
407 tf_device::ClusterFuncOp cluster_func, in HandleHostReplicatedInputs()
435 void HandleCluster(tf_device::ClusterFuncOp cluster_func, int32_t block_size, in HandleCluster()
605 Optional<tf_device::ClusterFuncOp> cluster_func; in runOnOperation() local
H A Dmark_input_output_aliases.cc54 tf_device::ClusterFuncOp cluster_func, in BuildAliasingInfo()
127 module.walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
H A Dannotate_parameter_replication.cc62 m.walk([&](tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dcluster_function_conversion.cc124 mlir::OpBuilder* builder, mlir::tf_device::ClusterFuncOp cluster_func) { in ReplaceClusterWithPartitionCallOp()
172 getOperation().walk([&](mlir::tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
177 for (auto cluster_func : llvm::reverse(clusters)) { in runOnOperation() local
H A Dset_default_sharding.cc59 getOperation().walk([&](mlir::tf_device::ClusterFuncOp cluster_func) { in runOnOperation()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/utils/
H A Dxla_sharding_util.cc206 mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder, in ExtractInputsForLogicalDevices()
307 mlir::tf_device::ClusterFuncOp cluster_func, in ParseAndValidateOutputSharding()
513 mlir::tf_device::ClusterFuncOp cluster_func, in GetOutputTypesForLogicalDeviceComputation()