| /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| #include "tensorflow/compiler/xla/service/mlir_gpu/passes.h" |
| |
| #include "absl/memory/memory.h" |
| #include "llvm/ADT/SetVector.h" |
| #include "mlir/Dialect/GPU/GPUDialect.h" // from @llvm-project |
| #include "mlir/Dialect/GPU/ParallelLoopMapper.h" // from @llvm-project |
| #include "mlir/Dialect/SCF/SCF.h" // from @llvm-project |
| #include "mlir/Dialect/SCF/Transforms.h" // from @llvm-project |
| #include "mlir/Dialect/StandardOps/IR/Ops.h" // from @llvm-project |
| #include "mlir/IR/BlockAndValueMapping.h" // from @llvm-project |
| #include "mlir/IR/Builders.h" // from @llvm-project |
| #include "mlir/Interfaces/SideEffectInterfaces.h" // from @llvm-project |
| #include "mlir/Transforms/LoopUtils.h" // from @llvm-project |
| #include "mlir/Transforms/RegionUtils.h" // from @llvm-project |
| #include "tensorflow/compiler/mlir/hlo/include/mlir-hlo/Dialect/mhlo/IR/lhlo_ops.h" |
| |
| namespace xla { |
| namespace mlir_gpu { |
| namespace { |
| |
| #define GEN_PASS_CLASSES |
| #include "tensorflow/compiler/xla/service/mlir_gpu/passes.h.inc" |
| |
| struct FusionOpRemoverPass : FusionOpRemoverPassBase<FusionOpRemoverPass> { |
| void runOnFunction() override { |
| getFunction().walk([&](mlir::lmhlo::FusionOp op) { |
| mlir::OpBuilder builder(op); |
| // FusionOp has a single region with a single block, so we can just walk |
| // over it and clone operations to the outside. |
| mlir::BlockAndValueMapping mapping; |
| for (auto& nested_op : op.region().front().without_terminator()) { |
| auto clone = builder.clone(nested_op, mapping); |
| for (auto pair : |
| llvm::zip(nested_op.getResults(), clone->getResults())) { |
| mapping.map(std::get<0>(pair), std::get<1>(pair)); |
| } |
| } |
| op.erase(); |
| }); |
| } |
| }; |
| |
| template <typename EffectTy> |
| bool HasEffectsOnValue(mlir::Value value, mlir::Operation* op) { |
| auto mem_effects_interface = |
| mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(op); |
| if (!mem_effects_interface) { |
| return false; |
| } |
| llvm::SmallVector<mlir::MemoryEffects::EffectInstance, 2> effects; |
| mem_effects_interface.getEffects(effects); |
| return llvm::any_of(effects, |
| [op](const mlir::MemoryEffects::EffectInstance& effect) { |
| return mlir::isa<EffectTy>(effect.getEffect()); |
| }); |
| } |
| |
| struct StoreForwardingPass : StoreForwardingPassBase<StoreForwardingPass> { |
| mlir::StoreOp findStore(mlir::Operation* op, |
| std::function<bool(mlir::StoreOp)> matches) { |
| // Search from op upwards in the current block. |
| mlir::Block* block = op->getBlock(); |
| auto startFromIt = |
| std::find_if(block->rbegin(), block->rend(), |
| [op](mlir::Operation& other) { return &other == op; }); |
| for (auto storeOpIt = startFromIt; storeOpIt != block->rend(); |
| ++storeOpIt) { |
| auto storeOp = llvm::dyn_cast<mlir::StoreOp>(&*(storeOpIt)); |
| if (!storeOp || !matches(storeOp)) { |
| continue; |
| } |
| |
| return storeOp; |
| } |
| // No store operation found. Continue search outside of the parallel |
| // loop if block is in a parallel loop. |
| if (auto parallelOp = |
| llvm::dyn_cast<mlir::scf::ParallelOp>(block->getParentOp())) { |
| return findStore(parallelOp.getOperation(), matches); |
| } |
| return {}; |
| } |
| |
| // Recursively search defining ops for AllocOp. Return either AllocOp if it is |
| // found or nullptr. |
| mlir::Operation* SearchAllocOp(mlir::Value memref) { |
| mlir::Operation* defOp = memref.getDefiningOp(); |
| while (auto subviewOp = mlir::dyn_cast_or_null<mlir::SubViewOp>(defOp)) { |
| defOp = subviewOp.source().getDefiningOp(); |
| } |
| return HasEffectsOnValue<mlir::MemoryEffects::Allocate>(memref, defOp) |
| ? defOp |
| : nullptr; |
| } |
| |
| // Retrieves AllocOp from the cache or actually looks for it. |
| mlir::Operation* GetAllocOp( |
| mlir::Value memref, |
| llvm::DenseMap<mlir::Value, mlir::Operation*>* memrefToAllocOp) { |
| auto allocOpIt = memrefToAllocOp->find(memref); |
| if (allocOpIt != memrefToAllocOp->end()) { |
| return allocOpIt->second; |
| } |
| mlir::Operation* allocOp = SearchAllocOp(memref); |
| memrefToAllocOp->insert({memref, allocOp}); |
| return allocOp; |
| } |
| |
| void runOnFunction() override { |
| llvm::DenseMap<mlir::Value, mlir::Operation*> memrefToAllocOp; |
| |
| getFunction().walk([&](mlir::LoadOp loadOp) { |
| auto storeOp = findStore(loadOp, [&](mlir::StoreOp storeOp) { |
| mlir::Operation* storeOpAlloc = |
| GetAllocOp(storeOp.memref(), &memrefToAllocOp); |
| mlir::Operation* loadOpAlloc = |
| GetAllocOp(loadOp.memref(), &memrefToAllocOp); |
| return storeOpAlloc && loadOpAlloc && (storeOpAlloc == loadOpAlloc); |
| }); |
| if (!storeOp) { |
| return; |
| } |
| auto storeIndices = storeOp.getIndices(); |
| auto loadIndices = loadOp.getIndices(); |
| if (!std::equal(storeIndices.begin(), storeIndices.end(), |
| loadIndices.begin(), loadIndices.end())) { |
| return; |
| } |
| loadOp.replaceAllUsesWith(storeOp.getValueToStore()); |
| loadOp.erase(); |
| }); |
| } |
| }; |
| |
| struct DeadTempBufferRemovalPass |
| : DeadTempBufferRemovalPassBase<DeadTempBufferRemovalPass> { |
| bool operationConsideredDead(mlir::Operation* op) { |
| for (auto result : op->getResults()) { |
| if (!llvm::all_of(result.getUsers(), [&](mlir::Operation* op) { |
| // Store and Dealloc is OK. |
| if (llvm::isa<mlir::StoreOp, mlir::DeallocOp>(op)) { |
| return true; |
| } |
| // Load without uses is also ok. |
| if (auto loadOp = llvm::dyn_cast<mlir::LoadOp>(op)) { |
| return loadOp.use_empty(); |
| } |
| // Subview is ok if it is dead itself. |
| if (llvm::isa<mlir::SubViewOp>(op)) { |
| return operationConsideredDead(op); |
| } |
| return false; |
| })) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| void recursiveErase(mlir::Operation* op, |
| llvm::SmallVectorImpl<mlir::Operation*>* erase_list) { |
| for (auto result : op->getResults()) { |
| for (auto user : llvm::make_early_inc_range(result.getUsers())) { |
| recursiveErase(user, erase_list); |
| } |
| } |
| erase_list->push_back(op); |
| } |
| |
| void runOnFunction() override { |
| llvm::SmallVector<mlir::Operation*, 8> dead_ops; |
| getFunction().walk([&](mlir::Operation* op) { |
| if (op->getNumResults() != 1 || |
| !HasEffectsOnValue<mlir::MemoryEffects::Allocate>(op->getResult(0), |
| op)) { |
| return; |
| } |
| if (!operationConsideredDead(op)) { |
| return; |
| } |
| |
| // TODO(herhut): There should be a generic helper for this. |
| recursiveErase(op, &dead_ops); |
| }); |
| for (auto op : dead_ops) { |
| op->erase(); |
| } |
| } |
| }; |
| |
| struct RewriteKernelSignaturePass |
| : RewriteKernelSignaturePassBase<RewriteKernelSignaturePass> { |
| void runOnFunction() override { |
| mlir::FuncOp func = getFunction(); |
| mlir::ModuleOp module = func.getParentOfType<mlir::ModuleOp>(); |
| getFunction().walk([&](mlir::gpu::LaunchFuncOp launchOp) { |
| mlir::gpu::GPUFuncOp kernel = |
| module.lookupSymbol<mlir::gpu::GPUFuncOp>(launchOp.kernel()); |
| |
| if (kernel.getNumFuncArguments() != |
| func.getNumArguments() + func.getNumResults()) { |
| kernel.emitError() |
| << "number of kernel arguments does not match number" |
| << "of arguments and results of surrounding function"; |
| signalPassFailure(); |
| return; |
| } |
| if (!llvm::hasSingleElement(func)) { |
| func.emitError() << "surrounding function has more than one block"; |
| signalPassFailure(); |
| return; |
| } |
| |
| // Compute a map from function arguments to kernel function operands. |
| mlir::BlockAndValueMapping func_to_kernel; |
| for (mlir::BlockArgument arg : func.getArguments()) { |
| for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) { |
| if (launchOp.getKernelOperand(i) == arg) { |
| func_to_kernel.map(arg, kernel.getArgument(i)); |
| break; |
| } |
| } |
| } |
| // Also add function results that are computed by the launch. |
| mlir::Operation* returnOp = func.getBody().back().getTerminator(); |
| for (mlir::Value result : returnOp->getOperands()) { |
| for (int i = 0, e = launchOp.getNumKernelOperands(); i < e; ++i) { |
| if (launchOp.getKernelOperand(i) == result) { |
| func_to_kernel.map(result, kernel.getArgument(i)); |
| break; |
| } |
| } |
| } |
| |
| // Create a new kernel function with modified signature. It will have the |
| // parameters and result types of the original funcion as its parameter |
| // type and otherwise will be void. |
| auto gpu_module = kernel.getParentOfType<mlir::gpu::GPUModuleOp>(); |
| mlir::OpBuilder kernel_builder(gpu_module.body()); |
| auto operand_types = llvm::to_vector<4>(llvm::concat<const mlir::Type>( |
| func.getType().getInputs(), func.getType().getResults())); |
| auto new_kernel = kernel_builder.create<mlir::gpu::GPUFuncOp>( |
| kernel.getLoc(), kernel.getName(), |
| kernel_builder.getFunctionType(operand_types, {})); |
| new_kernel->setAttr(mlir::gpu::GPUDialect::getKernelFuncAttrName(), |
| kernel_builder.getUnitAttr()); |
| |
| // Create a map from old kernel argument to new one. |
| mlir::BlockAndValueMapping old_kernel_to_new; |
| for (int i = 0, e = func.getNumArguments(); i < e; ++i) { |
| mlir::Value func_arg = func.getArgument(i); |
| mlir::Value new_kernel_arg = new_kernel.getArgument(i); |
| mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(func_arg); |
| if (!old_kernel_arg) { |
| kernel.emitOpError() |
| << "argument " << i |
| << " to containing function is not an argument to the kernel"; |
| signalPassFailure(); |
| return; |
| } |
| old_kernel_to_new.map(old_kernel_arg, new_kernel_arg); |
| } |
| for (int i = 0, e = returnOp->getNumOperands(); i < e; ++i) { |
| mlir::Value ret_op = returnOp->getOperand(i); |
| mlir::Value new_kernel_arg = |
| new_kernel.getArgument(func.getNumArguments() + i); |
| mlir::Value old_kernel_arg = func_to_kernel.lookupOrNull(ret_op); |
| if (!old_kernel_arg) { |
| kernel.emitOpError() |
| << "result " << i |
| << " of containing function is not an argument to the kernel"; |
| signalPassFailure(); |
| return; |
| } |
| old_kernel_to_new.map(old_kernel_arg, new_kernel_arg); |
| } |
| // Steal the body by appending the blocks and inserting a branch. |
| kernel.body().cloneInto(&new_kernel.getBody(), old_kernel_to_new); |
| kernel_builder.setInsertionPointToEnd(&new_kernel.body().front()); |
| kernel_builder.create<mlir::BranchOp>( |
| new_kernel.getLoc(), &*std::next(new_kernel.body().begin())); |
| // Now create a new launchOp calling the new kernel. We need to forward |
| // the arguments of the surrounding function and operands to the return. |
| mlir::SmallVector<mlir::Value, 4> new_operands; |
| new_operands.reserve(new_kernel.getNumFuncArguments()); |
| new_operands.append(func.args_begin(), func.args_end()); |
| new_operands.append(returnOp->operand_begin(), returnOp->operand_end()); |
| mlir::OpBuilder launch_builder(launchOp); |
| launch_builder.create<mlir::gpu::LaunchFuncOp>( |
| launchOp.getLoc(), new_kernel, launchOp.getGridSizeOperandValues(), |
| launchOp.getBlockSizeOperandValues(), new_operands); |
| // Launch does not have results, so we can just erase it. And the kernel |
| // also needs to go. |
| launchOp.erase(); |
| kernel.erase(); |
| }); |
| } |
| }; |
| |
| struct MapParallelLoopsPass : MapParallelLoopsPassBase<MapParallelLoopsPass> { |
| void runOnFunction() override { |
| mlir::greedilyMapParallelSCFToGPU(getFunction().getBody()); |
| } |
| }; |
| |
| struct FuseInnerParallelLoopsPass |
| : FuseInnerParallelLoopsPassBase<FuseInnerParallelLoopsPass> { |
| void runOnFunction() override { |
| getFunction().walk([](mlir::scf::ParallelOp op) { |
| mlir::scf::naivelyFuseParallelOps(op.region()); |
| }); |
| } |
| }; |
| |
| struct ParallelLoopCollapsingToFirstDimPass |
| : ParallelLoopCollapsingToFirstDimPassBase< |
| ParallelLoopCollapsingToFirstDimPass> { |
| void runOnFunction() override { |
| getFunction().walk([&](mlir::scf::ParallelOp op) { |
| unsigned num_loops = op.getNumLoops(); |
| std::vector<unsigned> combinedLoops; |
| combinedLoops.reserve(num_loops); |
| for (unsigned i = 0; i < num_loops; ++i) { |
| combinedLoops.push_back(i); |
| } |
| mlir::collapseParallelLoops(op, {combinedLoops}); |
| }); |
| } |
| }; |
| |
| } // namespace |
| |
| std::unique_ptr<mlir::FunctionPass> createFusionOpRemoverPass() { |
| return absl::make_unique<FusionOpRemoverPass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> createStoreForwardingPass() { |
| return absl::make_unique<StoreForwardingPass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> createDeadTempBufferRemovalPass() { |
| return absl::make_unique<DeadTempBufferRemovalPass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> createRewriteKernelSignaturePass() { |
| return absl::make_unique<RewriteKernelSignaturePass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> createFuseInnerParallelLoopsPass() { |
| return absl::make_unique<FuseInnerParallelLoopsPass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> createMapParallelLoopsPass() { |
| return absl::make_unique<MapParallelLoopsPass>(); |
| } |
| |
| std::unique_ptr<mlir::FunctionPass> |
| createParallelLoopCollapsingToFirstDimPass() { |
| return absl::make_unique<ParallelLoopCollapsingToFirstDimPass>(); |
| } |
| |
| } // namespace mlir_gpu |
| } // namespace xla |
