experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA:CPU] Register replica-id IR

Browse Source 
Otherwise any use of replica-id as an argument to another operation crashes.

PiperOrigin-RevId: 282934610
Change-Id: Iab3957f85910fd1f453c6d1b9043c9d220ece633
This commit is contained in:
Benjamin Kramer 2019-11-28 06:24:54 -08:00 committed by TensorFlower Gardener
parent d3645bd12b
commit d393702997
2 changed files with 3 additions and 1 deletions
tensorflow/compiler/xla
service/cpu
ir_emitter.cc
tests
collective_ops_test.cc

1
tensorflow/compiler/xla/service/cpu/ir_emitter.cc
View File

@ -1511,6 +1511,7 @@ Status IrEmitter::HandleAllReduce(HloInstruction* crs) {
}
Status IrEmitter::HandleReplicaId(HloInstruction* hlo) {
TF_RETURN_IF_ERROR(EmitTargetAddressForOp(hlo));
llvm::Type* i8_ptr_type = llvm::Type::getInt8PtrTy(module_->getContext());
llvm::FunctionType* replica_id_function_ty =
llvm::FunctionType::get(b_.getVoidTy(),

3
tensorflow/compiler/xla/tests/collective_ops_test.cc
View File

@ -397,7 +397,8 @@ XLA_TEST_F(CollectiveOpsTest, ReplicaId) {
const char* const kModuleStr = R"(
HloModule test
ENTRY test_computation {
ROOT id = u32[] replica-id()
id = u32[] replica-id()
ROOT out = u32[] copy(id)
}
)";
const int64 kNumReplicas = 4;