1 //===- TestLoopMapping.cpp --- Parametric loop mapping pass ---------------===// 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 a pass to parametrically map scf.for loops to virtual 10 // processing element dimensions. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Dialect/Affine/IR/AffineOps.h" 15 #include "mlir/Dialect/Affine/LoopUtils.h" 16 #include "mlir/Dialect/SCF/IR/SCF.h" 17 #include "mlir/IR/Builders.h" 18 #include "mlir/Pass/Pass.h" 19 20 using namespace mlir; 21 using namespace mlir::affine; 22 23 namespace { 24 struct TestLoopMappingPass 25 : public PassWrapper<TestLoopMappingPass, OperationPass<>> { MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID__anoncee15cae0111::TestLoopMappingPass26 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLoopMappingPass) 27 28 StringRef getArgument() const final { 29 return "test-mapping-to-processing-elements"; 30 } getDescription__anoncee15cae0111::TestLoopMappingPass31 StringRef getDescription() const final { 32 return "test mapping a single loop on a virtual processor grid"; 33 } 34 explicit TestLoopMappingPass() = default; 35 getDependentDialects__anoncee15cae0111::TestLoopMappingPass36 void getDependentDialects(DialectRegistry ®istry) const override { 37 registry.insert<affine::AffineDialect, scf::SCFDialect>(); 38 } 39 runOnOperation__anoncee15cae0111::TestLoopMappingPass40 void runOnOperation() override { 41 // SSA values for the transformation are created out of thin air by 42 // unregistered "new_processor_id_and_range" operations. This is enough to 43 // emulate mapping conditions. 44 SmallVector<Value, 8> processorIds, numProcessors; 45 getOperation()->walk([&processorIds, &numProcessors](Operation *op) { 46 if (op->getName().getStringRef() != "new_processor_id_and_range") 47 return; 48 processorIds.push_back(op->getResult(0)); 49 numProcessors.push_back(op->getResult(1)); 50 }); 51 52 getOperation()->walk([&processorIds, &numProcessors](scf::ForOp op) { 53 // Ignore nested loops. 54 if (op->getParentRegion()->getParentOfType<scf::ForOp>()) 55 return; 56 mapLoopToProcessorIds(op, processorIds, numProcessors); 57 }); 58 } 59 }; 60 } // namespace 61 62 namespace mlir { 63 namespace test { registerTestLoopMappingPass()64void registerTestLoopMappingPass() { PassRegistration<TestLoopMappingPass>(); } 65 } // namespace test 66 } // namespace mlir 67