}
};
+// Converts all_reduce op to LLVM/NVVM ops.
+struct GPUAllReduceOpLowering : public LLVMOpLowering {
+ explicit GPUAllReduceOpLowering(LLVMTypeConverter &lowering_)
+ : LLVMOpLowering(gpu::AllReduce::getOperationName(),
+ lowering_.getDialect()->getContext(), lowering_),
+ int32Type(LLVM::LLVMType::getInt32Ty(lowering_.getDialect())) {}
+
+ PatternMatchResult
+ matchAndRewrite(Operation *op, ArrayRef<Value *> operands,
+ ConversionPatternRewriter &rewriter) const override {
+ Value *result = createBlockReduce(op->getLoc(), operands.front(), rewriter);
+ rewriter.replaceOp(op, {result});
+ return matchSuccess();
+ }
+
+private:
+ // Creates an all_reduce across the local workgroup.
+ //
+ // First reduce the elements within a subgroup (i.e. warp). The first
+ // invocation of each subgroup writes the intermediate result to shared
+ // memory. After synchronizing the local workgroup, each subgroup reduces all
+ // values from shared memory.
+ //
+ // %warp_reduce = ... (see createWarpReduce)
+ // %buffer = llvm.mlir.addressof @reduce_buffer : !llvm<"[32 x float]*">
+ // %zero = llvm.mlir.constant(0 : i32) : !llvm.i32
+ // %lane_id = nvvm.read.ptx.sreg.laneid : !llvm.i32
+ // %is_first_lane = llvm.icmp "eq" %lane_id, %zero : !llvm.i32
+ // llvm.cond_br %is_first_lane, ^then, ^continue
+ // ^then:
+ // %warp_id = ... (see getWarpId)
+ // %store_dst = llvm.getelementptr %buffer[%zero, %warp_id]
+ // llvm.store %store_dst, %warp_reduce : !llvm.float
+ // llvm.br ^continue
+ // ^continue:
+ // nvvm.barrier0
+ // %load_src = llvm.getelementptr %buffer[%zero, %lane_id]
+ // %value = llvm.load %load_src : !llvm.float
+ // %result = ... (see createWarpReduce)
+ Value *createBlockReduce(Location loc, Value *operand,
+ ConversionPatternRewriter &rewriter) const {
+ auto type = operand->getType().cast<LLVM::LLVMType>();
+
+ Value *warpReduce = createWarpReduce(loc, operand, rewriter);
+
+ auto module = warpReduce->getDefiningOp()->getParentOfType<ModuleOp>();
+ assert(module && "op must belong to a module");
+ Value *sharedMemPtr =
+ createSharedMemoryArray(loc, module, type, kWarpSize, rewriter);
+
+ Value *zero = rewriter.create<LLVM::ConstantOp>(
+ loc, int32Type, rewriter.getI32IntegerAttr(0u));
+ Value *laneId = rewriter.create<NVVM::LaneIdOp>(loc, int32Type);
+ Value *isFirstLane = rewriter.create<LLVM::ICmpOp>(
+ loc, LLVM::ICmpPredicate::eq, laneId, zero);
+
+ Block *currentBlock = rewriter.getInsertionBlock();
+ auto currentPoint = rewriter.getInsertionPoint();
+
+ Block *thenBlock = rewriter.splitBlock(currentBlock, currentPoint);
+ Block *continueBlock = rewriter.splitBlock(thenBlock, currentPoint);
+
+ rewriter.setInsertionPointToEnd(currentBlock);
+ rewriter.create<LLVM::CondBrOp>(
+ loc, llvm::makeArrayRef(isFirstLane),
+ ArrayRef<Block *>{thenBlock, continueBlock});
+
+ rewriter.setInsertionPointToEnd(thenBlock);
+ Value *warpId = getWarpId(loc, rewriter);
+ Value *storeDst = rewriter.create<LLVM::GEPOp>(
+ loc, type, sharedMemPtr, ArrayRef<Value *>({zero, warpId}));
+ rewriter.create<LLVM::StoreOp>(loc, warpReduce, storeDst);
+ rewriter.create<LLVM::BrOp>(loc, ArrayRef<Value *>(),
+ llvm::makeArrayRef(continueBlock));
+
+ rewriter.setInsertionPointToStart(continueBlock);
+ rewriter.create<NVVM::Barrier0Op>(loc);
+ Value *loadSrc = rewriter.create<LLVM::GEPOp>(
+ loc, type, sharedMemPtr, ArrayRef<Value *>({zero, laneId}));
+ Value *value = rewriter.create<LLVM::LoadOp>(loc, type, loadSrc);
+ Value *result = createWarpReduce(loc, value, rewriter);
+
+ return result;
+ }
+
+ // Creates an all_reduce across the subgroup. Creates a preamble
+ //
+ // %active_mask = llvm.mlir.constant(-1 : i32) : !llvm.i32
+ // %mask_and_clamp = llvm.mlir.constant(31 : i32) : !llvm.i32
+ //
+ // plus the accumulation for i = 1, 2, 4, 8, 16:
+ //
+ // %offset = llvm.mlir.constant(i : i32) : !llvm.i32
+ // %value = nvvm.shfl.sync.bfly
+ // %active_mask, %operand, %offset, %mask_and_clamp : !llvm.float
+ // %operand = llvm.fadd %operand, %value : !llvm.float
+ //
+ // Each invocation returns the same result.
+ //
+ // Note: this currently only supports reducing exactly 32 values.
+ Value *createWarpReduce(Location loc, Value *operand,
+ ConversionPatternRewriter &rewriter) const {
+ // TODO(csigg): Generalize to partial warps and other types of accumulation.
+ static_assert(kWarpSize == 32, "Only warp size of 32 is supported.");
+ auto activeMask = rewriter.create<LLVM::ConstantOp>(
+ loc, int32Type, rewriter.getI32IntegerAttr(~0u));
+ auto maskAndClamp = rewriter.create<LLVM::ConstantOp>(
+ loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize - 1));
+
+ auto resultType = operand->getType();
+ for (int i = 1; i < kWarpSize; i <<= 1) {
+ auto offset = rewriter.create<LLVM::ConstantOp>(
+ loc, int32Type, rewriter.getI32IntegerAttr(i));
+ auto value = rewriter.create<NVVM::ShflBflyOp>(
+ loc, resultType, activeMask, operand, offset, maskAndClamp);
+ operand = rewriter.create<LLVM::FAddOp>(loc, resultType, operand, value);
+ }
+ return operand;
+ }
+
+ // Creates a global array stored in shared memory.
+ //
+ // llvm.mlir.global @reduce_buffer()
+ // {addr_space = 3 : i32} : !llvm<"[32 x float]">
+ //
+ Value *createSharedMemoryArray(Location loc, ModuleOp module,
+ LLVM::LLVMType elementType, int numElements,
+ ConversionPatternRewriter &rewriter) const {
+ OpBuilder builder(module.getBodyRegion());
+
+ auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
+ StringRef name = "reduce_buffer";
+ auto addrSpace =
+ builder.getNamedAttr("addr_space", builder.getI32IntegerAttr(3));
+ auto globalOp = builder.create<LLVM::GlobalOp>(
+ loc, arrayType.cast<LLVM::LLVMType>(),
+ /*isConstant=*/false, name, /*value=*/Attribute(),
+ llvm::makeArrayRef(addrSpace));
+
+ return rewriter.create<LLVM::AddressOfOp>(loc, globalOp);
+ }
+
+ // Returns the index of the subgroup within the local workgroup.
+ //
+ // %warp_size = llvm.mlir.constant(32 : i32) : !llvm.i32
+ // %thread_idx = nvvm.read.ptx.sreg.tid.x : !llvm.i32
+ // %warp_idx = llvm.sdiv %thread_idx, %warp_size : !llvm.i32
+ //
+ Value *getWarpId(Location loc, ConversionPatternRewriter &rewriter) const {
+ auto warpSize = rewriter.create<LLVM::ConstantOp>(
+ loc, int32Type, rewriter.getI32IntegerAttr(kWarpSize));
+ auto threadIdx = getLinearThreadIndex(loc, rewriter);
+ return rewriter.create<LLVM::SDivOp>(loc, int32Type, threadIdx, warpSize);
+ }
+
+ Value *getLinearThreadIndex(Location loc,
+ ConversionPatternRewriter &rewriter) const {
+ // TODO(csigg): support 2- and 3-dimensional blocks.
+ return rewriter.create<NVVM::ThreadIdXOp>(loc, int32Type);
+ }
+
+ LLVM::LLVMType int32Type;
+
+ // TODO(csigg): Support other warp sizes.
+ static constexpr int kWarpSize = 32;
+};
+
// A pass that replaces all occurences of GPU device operations with their
// corresponding NVVM equivalent.
//
GPUIndexIntrinsicOpLowering<gpu::BlockId, NVVM::BlockIdXOp,
NVVM::BlockIdYOp, NVVM::BlockIdZOp>,
GPUIndexIntrinsicOpLowering<gpu::GridDim, NVVM::GridDimXOp,
- NVVM::GridDimYOp, NVVM::GridDimZOp>>(
- converter);
+ NVVM::GridDimYOp, NVVM::GridDimZOp>,
+ GPUAllReduceOpLowering>(converter);
ConversionTarget target(getContext());
target.addLegalDialect<LLVM::LLVMDialect>();