1 //===- ParallelLoopMapper.cpp - Utilities for mapping parallel loops to GPU =// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements utilities to generate mappings for parallel loops to 10 // GPU devices. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/GPU/ParallelLoopMapper.h" 15 16 #include "mlir/Dialect/GPU/GPUDialect.h" 17 #include "mlir/Dialect/GPU/Passes.h" 18 #include "mlir/Dialect/LoopOps/LoopOps.h" 19 #include "mlir/IR/AffineMap.h" 20 #include "mlir/Pass/Pass.h" 21 22 using namespace mlir; 23 using namespace mlir::gpu; 24 using namespace mlir::loop; 25 26 namespace { 27 28 enum MappingLevel { MapGrid = 0, MapBlock = 1, Sequential = 2 }; 29 30 static constexpr int kNumHardwareIds = 3; 31 32 } // namespace 33 34 /// Bounded increment on MappingLevel. Increments to the next 35 /// level unless Sequential was already reached. 36 MappingLevel &operator++(MappingLevel &mappingLevel) { 37 if (mappingLevel < Sequential) { 38 mappingLevel = static_cast<MappingLevel>(mappingLevel + 1); 39 } 40 return mappingLevel; 41 } 42 43 /// Computed the hardware id to use for a given mapping level. Will 44 /// assign x,y and z hardware ids for the first 3 dimensions and use 45 /// sequential after. 46 static int64_t getHardwareIdForMapping(MappingLevel level, int dimension) { 47 if (dimension >= kNumHardwareIds || level == Sequential) 48 return Sequential * kNumHardwareIds; 49 return (level * kNumHardwareIds) + dimension; 50 } 51 52 /// Add mapping information to the given parallel loop. Do not add 53 /// mapping information if the loop already has it. Also, don't 54 /// start a mapping at a nested loop. 55 static void mapParallelOp(ParallelOp parallelOp, 56 MappingLevel mappingLevel = MapGrid) { 57 // Do not try to add a mapping to already mapped loops or nested loops. 58 if (parallelOp.getAttr(gpu::kMappingAttributeName) || 59 ((mappingLevel == MapGrid) && parallelOp.getParentOfType<ParallelOp>())) 60 return; 61 62 MLIRContext *ctx = parallelOp.getContext(); 63 Builder b(ctx); 64 SmallVector<Attribute, 4> attrs; 65 attrs.reserve(parallelOp.getNumInductionVars()); 66 for (int i = 0, e = parallelOp.getNumInductionVars(); i < e; ++i) { 67 SmallVector<NamedAttribute, 3> entries; 68 entries.emplace_back(b.getNamedAttr( 69 kProcessorEntryName, 70 b.getI64IntegerAttr(getHardwareIdForMapping(mappingLevel, i)))); 71 entries.emplace_back(b.getNamedAttr( 72 kIndexMapEntryName, AffineMapAttr::get(b.getDimIdentityMap()))); 73 entries.emplace_back(b.getNamedAttr( 74 kBoundMapEntryName, AffineMapAttr::get(b.getDimIdentityMap()))); 75 attrs.push_back(DictionaryAttr::get(entries, ctx)); 76 } 77 parallelOp.setAttr(kMappingAttributeName, ArrayAttr::get(attrs, ctx)); 78 ++mappingLevel; 79 // Parallel loop operations are immediately nested, so do not use 80 // walk but just iterate over the operations. 81 for (Operation &op : *parallelOp.getBody()) { 82 if (ParallelOp nested = dyn_cast<ParallelOp>(op)) 83 mapParallelOp(nested, mappingLevel); 84 } 85 } 86 87 void mlir::greedilyMapParallelLoopsToGPU(Region ®ion) { 88 region.walk([](ParallelOp parallelOp) { mapParallelOp(parallelOp); }); 89 } 90