Home
last modified time | relevance | path

Searched refs:function_block (Results 1 – 2 of 2) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtpu_sharding_identification_pass.cc114 Block& function_block = func.front(); in VerifyShardings() local
116 llvm::zip(sharding_for_args, function_block.getArguments())) { in VerifyShardings()
121 Operation* terminator = function_block.getTerminator(); in VerifyShardings()
187 Block& function_block = func.front(); in IdentifyXlaShardingForComputationInputs() local
189 sharding_for_args.reserve(function_block.getNumArguments()); in IdentifyXlaShardingForComputationInputs()
209 llvm::zip(cluster_func.operands(), function_block.getArguments())) { in IdentifyXlaShardingForComputationInputs()
357 Block& function_block = func.front(); in IdentifyXlaShardingForComputationOutputs() local
358 Operation* terminator = function_block.getTerminator(); in IdentifyXlaShardingForComputationOutputs()
471 Block& function_block = func.front(); in IdentifyXlaShardingForTPUComputation() local
473 llvm::zip(sharding_for_args, function_block.getArguments())) { in IdentifyXlaShardingForTPUComputation()
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/dtensor/mlir/
H A Dtpu_integration.cc106 auto& function_block = function->getCallableRegion()->front(); in CreateTPUCluster() local
107 builder->setInsertionPointToStart(&function_block); in CreateTPUCluster()
113 auto& function_body = function_block.getOperations(); in CreateTPUCluster()
119 mlir::Operation* function_block_terminator = function_block.getTerminator(); in CreateTPUCluster()