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/SCF.h"
17 #include "mlir/IR/Builders.h"
18 #include "mlir/Pass/Pass.h"
19 
20 #include "llvm/ADT/SetVector.h"
21 
22 using namespace mlir;
23 
24 namespace {
25 class TestLoopMappingPass
26     : public PassWrapper<TestLoopMappingPass, OperationPass<FuncOp>> {
27 public:
28   StringRef getArgument() const final {
29     return "test-mapping-to-processing-elements";
30   }
31   StringRef getDescription() const final {
32     return "test mapping a single loop on a virtual processor grid";
33   }
34   explicit TestLoopMappingPass() = default;
35 
36   void getDependentDialects(DialectRegistry &registry) const override {
37     registry.insert<AffineDialect, scf::SCFDialect>();
38   }
39 
40   void runOnOperation() override {
41     FuncOp func = getOperation();
42 
43     // SSA values for the transformation are created out of thin air by
44     // unregistered "new_processor_id_and_range" operations. This is enough to
45     // emulate mapping conditions.
46     SmallVector<Value, 8> processorIds, numProcessors;
47     func.walk([&processorIds, &numProcessors](Operation *op) {
48       if (op->getName().getStringRef() != "new_processor_id_and_range")
49         return;
50       processorIds.push_back(op->getResult(0));
51       numProcessors.push_back(op->getResult(1));
52     });
53 
54     func.walk([&processorIds, &numProcessors](scf::ForOp op) {
55       // Ignore nested loops.
56       if (op->getParentRegion()->getParentOfType<scf::ForOp>())
57         return;
58       mapLoopToProcessorIds(op, processorIds, numProcessors);
59     });
60   }
61 };
62 } // namespace
63 
64 namespace mlir {
65 namespace test {
66 void registerTestLoopMappingPass() { PassRegistration<TestLoopMappingPass>(); }
67 } // namespace test
68 } // namespace mlir
69