tensorflow/compiler/xla/service/gpu/ir_emitter.cc
// Implements atomic binary operations using atomic compare-and-swap
// (atomicCAS) as follows:
//   1. Reads the value from the memory pointed to by output_address and
//     records it as old_output.
//   2. Uses old_output as one of the source operand to perform the binary
//     operation and stores the result in new_output.
//   3. Calls atomicCAS which implements compare-and-swap as an atomic
//     operation. In particular, atomicCAS reads the value from the memory
//     pointed to by output_address, and compares the value with old_output. If
//     the two values equal, new_output is written to the same memory location
//     and true is returned to indicate that the atomic operation succeeds.
//     Otherwise, the new value read from the memory is returned. In this case,
//     the new value is copied to old_output, and steps 2. and 3. are repeated
//     until atomicCAS succeeds.
//
// On Nvidia GPUs, atomicCAS can only operate on 32 bit and 64 bit integers. If
// the element type of the binary operation is 32 bits or 64 bits, the integer
// type of the same size is used for the atomicCAS operation. On the other hand,
// if the element type is smaller than 32 bits, int32_t is used for the
// atomicCAS operation. In this case, atomicCAS reads and writes 32 bit values
// from the memory, which is larger than the memory size required by the
// original atomic binary operation. We mask off the last two bits of the
// output_address and use the result as an address to read the 32 bit values
// from the memory. This can avoid out of bound memory accesses if tensor
// buffers are 4 byte aligned and have a size of 4N, an assumption that the
// runtime can guarantee.
//
// The pseudo code is shown below. Variables *_address are pointers to a memory
// region with a size equal to the size of the atomicCAS operation, with the
// exception that new_output_address is a pointer to a memory region with a size
// equal to the element size of the binary operation.
//
//   element_size = sizeof(element_type);
//   atomic_size = max(32, element_size);
//   cas_new_output_address = alloca(atomic_size);
//   cas_old_output_address = alloca(atomic_size);
//   if (atomic_size != element_size) {
//     atomic_address = output_address & ((int64_t)(-4));
//     new_output_address = cas_new_output_address + (output_address & 3);
//   } else {
//     atomic_address = output_address;
//     new_output_address = cas_new_output_address;
//   }
//
//   *cas_old_output_address = *atomic_address;
//   do {
//     *cas_new_output_address = *cas_old_output_address;
//     *new_output_address = operation(*new_output_address, *source_address);
//     (*cas_old_output_address, success) =
//       atomicCAS(atomic_address, *cas_old_output_address,
//       *cas_new_output_address);
//   } while (!success);
//
Implements atomic binary operations using atomic compare-and-swap
(atomicCAS) as follows:
1. Reads the value from the memory pointed to by output_address and
records it as old_output.
2. Uses old_output as one of the source operand to perform the binary
operation and stores the result in new_output.
3. Calls atomicCAS which implements compare-and-swap as an atomic
operation. In particular, atomicCAS reads the value from the memory
pointed to by output_address, and compares the value with old_output. If
the two values equal, new_output is written to the same memory location
and true is returned to indicate that the atomic operation succeeds.
Otherwise, the new value read from the memory is returned. In this case,
the new value is copied to old_output, and steps 2. and 3. are repeated
until atomicCAS succeeds.

On Nvidia GPUs, atomicCAS can only operate on 32 bit and 64 bit integers. If
the element type of the binary operation is 32 bits or 64 bits, the integer
type of the same size is used for the atomicCAS operation. On the other hand,
if the element type is smaller than 32 bits, int32_t is used for the
atomicCAS operation. In this case, atomicCAS reads and writes 32 bit values
from the memory, which is larger than the memory size required by the
original atomic binary operation. We mask off the last two bits of the
output_address and use the result as an address to read the 32 bit values
from the memory. This can avoid out of bound memory accesses if tensor
buffers are 4 byte aligned and have a size of 4N, an assumption that the
runtime can guarantee.

The pseudo code is shown below. Variables *_address are pointers to a memory
region with a size equal to the size of the atomicCAS operation, with the
exception that new_output_address is a pointer to a memory region with a size
equal to the element size of the binary operation.




Implements atomic binary operations using atomic compare - and - swap ( atomicCAS ) as follows : 1 . Reads the value from the memory pointed to by output_address and records it as old_output . 2 . Uses old_output as one of the source operand to perform the binary operation and stores the result in new_output . 3 . Calls atomicCAS which implements compare - and - swap as an atomic operation . In particular atomicCAS reads the value from the memory pointed to by output_address and compares the value with old_output . If the two values equal new_output is written to the same memory location and true is returned to indicate that the atomic operation succeeds . Otherwise the new value read from the memory is returned . In this case the new value is copied to old_output and steps 2 . and 3 . are repeated until atomicCAS succeeds . On Nvidia GPUs atomicCAS can only operate on 32 bit and 64 bit integers . If the element type of the binary operation is 32 bits or 64 bits the integer type of the same size is used for the atomicCAS operation . On the other hand if the element type is smaller than 32 bits int32_t is used for the atomicCAS operation . In this case atomicCAS reads and writes 32 bit values from the memory which is larger than the memory size required by the original atomic binary operation . We mask off the last two bits of the output_address and use the result as an address to read the 32 bit values from the memory . This can avoid out of bound memory accesses if tensor buffers are 4 byte aligned and have a size of 4N an assumption that the runtime can guarantee . The pseudo code is shown below . Variables * _address are pointers to a memory region with a size equal to the size of the atomicCAS operation with the exception that new_output_address is a pointer to a memory region with a size equal to the element size of the binary operation .
Status IrEmitter::EmitAtomicOperationUsingCAS(const HloComputation& computation,
                                              llvm::Value* output_address,
                                              llvm::Value* source_address) {
  llvm::PointerType* output_address_type =
      llvm::dyn_cast<llvm::PointerType>(output_address->getType());
  CHECK_NE(output_address_type, nullptr);
  llvm::Type* element_type = output_address_type->getPointerElementType();
  int element_size = llvm_ir::GetSizeInBits(element_type);
  int atomic_size = (element_size < 32) ? 32 : element_size;
  llvm::Type* atomic_type = b_.getIntNTy(atomic_size);
  llvm::Type* atomic_address_type =
      atomic_type->getPointerTo(output_address_type->getPointerAddressSpace());
  llvm::Value* cas_old_output_address = llvm_ir::EmitAllocaAtFunctionEntry(
      atomic_type, "cas_old_output_address", &b_);
  llvm::Value* cas_new_output_address = llvm_ir::EmitAllocaAtFunctionEntry(
      atomic_type, "cas_new_output_address", &b_);
  llvm::BasicBlock* loop_preheader_bb = b_.GetInsertBlock();
  llvm::Value* atomic_memory_address;
  llvm::Value* binop_output_address;
  if (element_size < 32) {
    CHECK_EQ((element_size % sizeof(char)), 0);
    llvm::Type* address_int_type =
        module_->getDataLayout().getIntPtrType(output_address_type);
    atomic_memory_address = PtrToInt(output_address, address_int_type);
    llvm::Value* mask = llvm::ConstantInt::get(address_int_type, 3);
    llvm::Value* offset = And(atomic_memory_address, mask);
    mask = llvm::ConstantInt::get(address_int_type, -4);
    atomic_memory_address = And(atomic_memory_address, mask);
    atomic_memory_address =
        IntToPtr(atomic_memory_address, atomic_address_type);
    binop_output_address =
        Add(PtrToInt(cas_new_output_address, address_int_type), offset);
    binop_output_address = IntToPtr(
        binop_output_address,
        llvm::PointerType::get(
            element_type,
            cas_new_output_address->getType()->getPointerAddressSpace()));
  } else {
    atomic_memory_address = b_.CreatePointerBitCastOrAddrSpaceCast(
        output_address, atomic_address_type);
    binop_output_address = b_.CreatePointerBitCastOrAddrSpaceCast(
        cas_new_output_address,
        llvm::PointerType::get(
            element_type,
            cas_new_output_address->getType()->getPointerAddressSpace()));
  }
  llvm::Value* cas_old_output = Load(atomic_memory_address, "cas_old_output");
  Store(cas_old_output, cas_old_output_address);
  llvm::BasicBlock* loop_exit_bb = loop_preheader_bb->splitBasicBlock(
      b_.GetInsertPoint(), "atomic_op_loop_exit");
  llvm::BasicBlock* loop_body_bb = llvm::BasicBlock::Create(
      b_.getContext(), "atomic_op_loop_body", b_.GetInsertBlock()->getParent());
  b_.SetInsertPoint(loop_body_bb);
  loop_preheader_bb->getTerminator()->setSuccessor(0, loop_body_bb);
   Use cas_old_output to initialize cas_new_output.
  cas_old_output = Load(cas_old_output_address, "cas_old_output");
  Store(cas_old_output, cas_new_output_address);
   Emits code to calculate new_output = operation(old_output, source);
  TF_RETURN_IF_ERROR(EmitCallToNestedComputation(
      computation, {binop_output_address, source_address},
      binop_output_address));
  llvm::Value* cas_new_output = Load(cas_new_output_address, "cas_new_output");
   If cas_new_output == cas_old_output, we're not asking for anything to
   change, so we're done here!
  llvm::Value* old_eq_new = ICmpEQ(cas_old_output, cas_new_output);
  llvm::BasicBlock* loop_cas_bb = llvm::BasicBlock::Create(
      b_.getContext(), "atomic_op_loop_cas", b_.GetInsertBlock()->getParent());
  CondBr(old_eq_new, loop_exit_bb, loop_cas_bb);
  b_.SetInsertPoint(loop_cas_bb);
   Emit code to perform the atomicCAS operation
   (cas_old_output, success) = atomicCAS(memory_address, cas_old_output,
                                         cas_new_output);
  llvm::Value* ret_value = AtomicCmpXchg(
      atomic_memory_address, cas_old_output, cas_new_output, llvm::MaybeAlign(),
      llvm::AtomicOrdering::SequentiallyConsistent,
      llvm::AtomicOrdering::SequentiallyConsistent, DetermineSyncScope());
   Extract the memory value returned from atomicCAS and store it as
  Store(ExtractValue(ret_value, 0, "cas_old_output"), cas_old_output_address);
   Extract the success bit returned from atomicCAS and generate a
   conditional branch on the success bit.
  CondBr(ExtractValue(ret_value, 1, "success"), loop_exit_bb, loop_body_bb);
   Set the insertion point to the exit basic block so that the caller of
   this method can continue emitting code to the right place.
  SetToFirstInsertPoint(loop_exit_bb, &b_);
  return Status::OK();
}
Status IrEmitter :: EmitAtomicOperationUsingCAS ( const HloComputation & computation , llvm :: Value * output_address , llvm :: Value * source_address ) { llvm :: PointerType * output_address_type = llvm :: dyn_cast < llvm :: PointerType > ( output_address -> getType ( ) ) ; CHECK_NE ( output_address_type , nullptr ) ; llvm :: Type * element_type = output_address_type -> getPointerElementType ( ) ; int element_size = llvm_ir :: GetSizeInBits ( element_type ) ; int atomic_size = ( element_size < 32 ) ? 32 : element_size ; llvm :: Type * atomic_type = b_ . getIntNTy ( atomic_size ) ; llvm :: Type * atomic_address_type = atomic_type -> getPointerTo ( output_address_type -> getPointerAddressSpace ( ) ) ; llvm :: Value * cas_old_output_address = llvm_ir :: EmitAllocaAtFunctionEntry ( atomic_type , " " , & b_ ) ; llvm :: Value * cas_new_output_address = llvm_ir :: EmitAllocaAtFunctionEntry ( atomic_type , " " , & b_ ) ; llvm :: BasicBlock * loop_preheader_bb = b_ . GetInsertBlock ( ) ; llvm :: Value * atomic_memory_address ; llvm :: Value * binop_output_address ; if ( element_size < 32 ) { CHECK_EQ ( ( element_size % sizeof ( char ) ) , 0 ) ; llvm :: Type * address_int_type = module_ -> getDataLayout ( ) . getIntPtrType ( output_address_type ) ; atomic_memory_address = PtrToInt ( output_address , address_int_type ) ; llvm :: Value * mask = llvm :: ConstantInt :: get ( address_int_type , 3 ) ; llvm :: Value * offset = And ( atomic_memory_address , mask ) ; mask = llvm :: ConstantInt :: get ( address_int_type , -4 ) ; atomic_memory_address = And ( atomic_memory_address , mask ) ; atomic_memory_address = IntToPtr ( atomic_memory_address , atomic_address_type ) ; binop_output_address = Add ( PtrToInt ( cas_new_output_address , address_int_type ) , offset ) ; binop_output_address = IntToPtr ( binop_output_address , llvm :: PointerType :: get ( element_type , cas_new_output_address -> getType ( ) -> getPointerAddressSpace ( ) ) ) ; } else { atomic_memory_address = b_ . CreatePointerBitCastOrAddrSpaceCast ( output_address , atomic_address_type ) ; binop_output_address = b_ . CreatePointerBitCastOrAddrSpaceCast ( cas_new_output_address , llvm :: PointerType :: get ( element_type , cas_new_output_address -> getType ( ) -> getPointerAddressSpace ( ) ) ) ; } llvm :: Value * cas_old_output = Load ( atomic_memory_address , " " ) ; Store ( cas_old_output , cas_old_output_address ) ; llvm :: BasicBlock * loop_exit_bb = loop_preheader_bb -> splitBasicBlock ( b_ . GetInsertPoint ( ) , " " ) ; llvm :: BasicBlock * loop_body_bb = llvm :: BasicBlock :: Create ( b_ . getContext ( ) , " " , b_ . GetInsertBlock ( ) -> getParent ( ) ) ; b_ . SetInsertPoint ( loop_body_bb ) ; loop_preheader_bb -> getTerminator ( ) -> setSuccessor ( 0 , loop_body_bb ) ; cas_old_output = Load ( cas_old_output_address , " " ) ; Store ( cas_old_output , cas_new_output_address ) ; TF_RETURN_IF_ERROR ( EmitCallToNestedComputation ( computation , { binop_output_address , source_address } , binop_output_address ) ) ; llvm :: Value * cas_new_output = Load ( cas_new_output_address , " " ) ; llvm :: Value * old_eq_new = ICmpEQ ( cas_old_output , cas_new_output ) ; llvm :: BasicBlock * loop_cas_bb = llvm :: BasicBlock :: Create ( b_ . getContext ( ) , " " , b_ . GetInsertBlock ( ) -> getParent ( ) ) ; CondBr ( old_eq_new , loop_exit_bb , loop_cas_bb ) ; b_ . SetInsertPoint ( loop_cas_bb ) ; llvm :: Value * ret_value = AtomicCmpXchg ( atomic_memory_address , cas_old_output , cas_new_output , llvm :: MaybeAlign ( ) , llvm :: AtomicOrdering :: SequentiallyConsistent , llvm :: AtomicOrdering :: SequentiallyConsistent , DetermineSyncScope ( ) ) ; Store ( ExtractValue ( ret_value , 0 , " " ) , cas_old_output_address ) ; CondBr ( ExtractValue ( ret_value , 1 , " " ) , loop_exit_bb , loop_body_bb ) ; SetToFirstInsertPoint ( loop_exit_bb , & b_ ) ; return Status :: OK ( ) ; }
Implements atomic binary operations using atomic compare-and-swap
(atomicCAS) as follows:
1.
Implements atomic binary operations using atomic compare - and - swap ( atomicCAS ) as follows : 1 .
// element_type is the data type for the binary operation. // cas_old_output_address and cas_new_output_address point to the scratch // memory where we store the old and new values for the repeated atomicCAS // operations. // Emit preparation code to the preheader. // binop_output_address points to the scratch memory that stores the // result of the binary operation. // Assume the element size is an integer number of bytes. // Use the value from the memory that atomicCAS operates on to initialize // cas_old_output. // Change preheader's successor from loop_exit_bb to loop_body_bb. // Emit the body of the loop that repeatedly invokes atomicCAS. // // Use cas_old_output to initialize cas_new_output. // Emits code to calculate new_output = operation(old_output, source); // If cas_new_output == cas_old_output, we're not asking for anything to // change, so we're done here! // Emit code to perform the atomicCAS operation // (cas_old_output, success) = atomicCAS(memory_address, cas_old_output, //                                       cas_new_output); // Extract the memory value returned from atomicCAS and store it as // cas_old_output. // Extract the success bit returned from atomicCAS and generate a // conditional branch on the success bit. // Set the insertion point to the exit basic block so that the caller of // this method can continue emitting code to the right place.