[Mlir-commits] [mlir] 0736bbd - [mlir][scf] Add callback to annotate ops during pipelining
Thomas Raoux
llvmlistbot at llvm.org
Tue Feb 15 12:48:28 PST 2022
Author: Thomas Raoux
Date: 2022-02-15T12:48:01-08:00
New Revision: 0736bbd7e2f7b43a0246945aaff851cff4827682
URL: https://github.com/llvm/llvm-project/commit/0736bbd7e2f7b43a0246945aaff851cff4827682
DIFF: https://github.com/llvm/llvm-project/commit/0736bbd7e2f7b43a0246945aaff851cff4827682.diff
LOG: [mlir][scf] Add callback to annotate ops during pipelining
This allow user to register a callback that can annotate operations
during software pipelining. This allows user potential annotate op to
know what part of the pipeline they correspond to.
Differential Revision: https://reviews.llvm.org/D119866
Added:
Modified:
mlir/include/mlir/Dialect/SCF/Transforms.h
mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
mlir/test/Dialect/SCF/loop-pipelining.mlir
mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/SCF/Transforms.h b/mlir/include/mlir/Dialect/SCF/Transforms.h
index 391812641c685..a66284562b765 100644
--- a/mlir/include/mlir/Dialect/SCF/Transforms.h
+++ b/mlir/include/mlir/Dialect/SCF/Transforms.h
@@ -125,7 +125,21 @@ struct PipeliningOption {
/// order picked for the pipelined loop.
using GetScheduleFnType = std::function<void(
scf::ForOp, std::vector<std::pair<Operation *, unsigned>> &)>;
- GetScheduleFnType getScheduleFn;
+ GetScheduleFnType getScheduleFn = nullptr;
+ enum class PipelinerPart {
+ Prologue,
+ Kernel,
+ Epilogue,
+ };
+ /// Lambda called by the pipeliner to allow the user to annotate the IR while
+ /// it is generated.
+ /// The callback passes the operation created along with the part of the
+ /// pipeline and the iteration index. The iteration index is always 0 for the
+ /// kernel. For the prologue and epilogue, it corresponds to the iteration
+ /// peeled out of the loop in the range [0, maxStage[.
+ using AnnotationlFnType =
+ std::function<void(Operation *, PipelinerPart, unsigned)>;
+ AnnotationlFnType annotateFn = nullptr;
// TODO: add option to decide if the prologue/epilogue should be peeled.
};
diff --git a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
index 8e20906e72514..493a1b891e343 100644
--- a/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/LoopPipelining.cpp
@@ -41,6 +41,7 @@ struct LoopPipelinerInternal {
int64_t ub;
int64_t lb;
int64_t step;
+ PipeliningOption::AnnotationlFnType annotateFn = nullptr;
// When peeling the kernel we generate several version of each value for
//
diff erent stage of the prologue. This map tracks the mapping between
@@ -126,6 +127,7 @@ bool LoopPipelinerInternal::initializeLoopInfo(
return !def || stages.find(def) == stages.end();
}))
return false;
+ annotateFn = options.annotateFn;
return true;
}
@@ -150,6 +152,8 @@ void LoopPipelinerInternal::emitPrologue(PatternRewriter &rewriter) {
if (it != valueMapping.end())
newOp->setOperand(opIdx, it->second[i - stages[op]]);
}
+ if (annotateFn)
+ annotateFn(newOp, PipeliningOption::PipelinerPart::Prologue, i);
for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
setValueMapping(op->getResult(destId), newOp->getResult(destId),
i - stages[op]);
@@ -297,6 +301,8 @@ void LoopPipelinerInternal::createKernel(
newOp->setOperand(operand.getOperandNumber(),
newForOp.getRegionIterArgs()[remap->second]);
}
+ if (annotateFn)
+ annotateFn(newOp, PipeliningOption::PipelinerPart::Kernel, 0);
}
// Collect the Values that need to be returned by the forOp. For each
@@ -363,6 +369,8 @@ LoopPipelinerInternal::emitEpilogue(PatternRewriter &rewriter) {
newOp->setOperand(opIdx, v);
}
}
+ if (annotateFn)
+ annotateFn(newOp, PipeliningOption::PipelinerPart::Epilogue, i - 1);
for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
setValueMapping(op->getResult(destId), newOp->getResult(destId),
maxStage - stages[op] + i);
diff --git a/mlir/test/Dialect/SCF/loop-pipelining.mlir b/mlir/test/Dialect/SCF/loop-pipelining.mlir
index 545a0ce981f4e..e2c740d1e56a2 100644
--- a/mlir/test/Dialect/SCF/loop-pipelining.mlir
+++ b/mlir/test/Dialect/SCF/loop-pipelining.mlir
@@ -1,4 +1,5 @@
// RUN: mlir-opt %s -test-scf-pipelining -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -test-scf-pipelining=annotate -split-input-file | FileCheck %s --check-prefix ANNOTATE
// CHECK-LABEL: simple_pipeline(
// CHECK-SAME: %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
@@ -97,6 +98,22 @@ func @simple_pipeline_step(%A: memref<?xf32>, %result: memref<?xf32>) {
// CHECK-NEXT: memref.store %[[LR]]#0, %[[R]][%[[C2]]] : memref<?xf32>
// CHECK-NEXT: %[[ADD2:.*]] = arith.addf %[[LR]]#1, %{{.*}} : f32
// CHECK-NEXT: memref.store %[[ADD2]], %[[R]][%[[C3]]] : memref<?xf32>
+
+// Prologue:
+// ANNOTATE: memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "prologue"}
+// ANNOTATE: memref.load {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "prologue"}
+// Kernel:
+// ANNOTATE: scf.for
+// ANNOTATE: memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
+// ANNOTATE: arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
+// ANNOTATE: memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
+// ANNOTATE: scf.yield
+// ANNOTATE: }
+// Epilogue:
+// ANNOTATE: memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
+// ANNOTATE: arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
+// ANNOTATE: memref.store {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "epilogue"}
+
func @three_stage(%A: memref<?xf32>, %result: memref<?xf32>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
diff --git a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
index f3106ed170807..0feb3e04bc717 100644
--- a/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
+++ b/mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp
@@ -90,12 +90,23 @@ static const StringLiteral kTestPipeliningStageMarker =
static const StringLiteral kTestPipeliningOpOrderMarker =
"__test_pipelining_op_order__";
+static const StringLiteral kTestPipeliningAnnotationPart =
+ "__test_pipelining_part";
+static const StringLiteral kTestPipeliningAnnotationIteration =
+ "__test_pipelining_iteration";
+
class TestSCFPipeliningPass
: public PassWrapper<TestSCFPipeliningPass, OperationPass<FuncOp>> {
public:
+ TestSCFPipeliningPass() = default;
+ TestSCFPipeliningPass(const TestSCFPipeliningPass &) {}
StringRef getArgument() const final { return "test-scf-pipelining"; }
StringRef getDescription() const final { return "test scf.forOp pipelining"; }
- explicit TestSCFPipeliningPass() = default;
+
+ Option<bool> annotatePipeline{
+ *this, "annotate",
+ llvm::cl::desc("Annote operations during loop pipelining transformation"),
+ llvm::cl::init(false)};
static void
getSchedule(scf::ForOp forOp,
@@ -115,6 +126,25 @@ class TestSCFPipeliningPass
});
}
+ static void annotate(Operation *op,
+ mlir::scf::PipeliningOption::PipelinerPart part,
+ unsigned iteration) {
+ OpBuilder b(op);
+ switch (part) {
+ case mlir::scf::PipeliningOption::PipelinerPart::Prologue:
+ op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("prologue"));
+ break;
+ case mlir::scf::PipeliningOption::PipelinerPart::Kernel:
+ op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("kernel"));
+ break;
+ case mlir::scf::PipeliningOption::PipelinerPart::Epilogue:
+ op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("epilogue"));
+ break;
+ }
+ op->setAttr(kTestPipeliningAnnotationIteration,
+ b.getI32IntegerAttr(iteration));
+ }
+
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<arith::ArithmeticDialect, StandardOpsDialect>();
}
@@ -123,7 +153,8 @@ class TestSCFPipeliningPass
RewritePatternSet patterns(&getContext());
mlir::scf::PipeliningOption options;
options.getScheduleFn = getSchedule;
-
+ if (annotatePipeline)
+ options.annotateFn = annotate;
scf::populateSCFLoopPipeliningPatterns(patterns, options);
(void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
getOperation().walk([](Operation *op) {
More information about the Mlir-commits
mailing list