diff --git a/mlir/compiler-course/sozonov_i_trace_loop_iterations/CMakeLists.txt b/mlir/compiler-course/sozonov_i_trace_loop_iterations/CMakeLists.txt new file mode 100644 index 0000000000000..e2d1a0ed20286 --- /dev/null +++ b/mlir/compiler-course/sozonov_i_trace_loop_iterations/CMakeLists.txt @@ -0,0 +1,16 @@ +set(Title "TraceLoopIterationsPass") +set(Student "Sozonov_Ilya") +set(Group "FIIT3") +set(TARGET_NAME "${Title}_${Student}_${Group}_MLIR") + +file(GLOB_RECURSE SOURCES *.cpp *.h *.hpp) + +add_llvm_pass_plugin(${TARGET_NAME} + ${SOURCES} + DEPENDS + intrinsics_gen + MLIRBuiltinLocationAttributesIncGen + BUILDTREE_ONLY +) + +set(MLIR_TEST_DEPENDS ${TARGET_NAME} ${MLIR_TEST_DEPENDS} PARENT_SCOPE) diff --git a/mlir/compiler-course/sozonov_i_trace_loop_iterations/TraceLoopIterationsPass.cpp b/mlir/compiler-course/sozonov_i_trace_loop_iterations/TraceLoopIterationsPass.cpp new file mode 100644 index 0000000000000..e04bdbb88588c --- /dev/null +++ b/mlir/compiler-course/sozonov_i_trace_loop_iterations/TraceLoopIterationsPass.cpp @@ -0,0 +1,88 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Tools/Plugins/PassPlugin.h" +#include "llvm/Support/raw_ostream.h" + +using namespace mlir; + +namespace { +class TraceLoopIterationsPass + : public PassWrapper> { +public: + StringRef getArgument() const final { + return "TraceLoopIterationsPass_Sozonov_Ilya_FIIT3_MLIR"; + } + StringRef getDescription() const final { + return "Insert trace calls at beginning and end of each loop iteration"; + } + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + OpBuilder builder(moduleOp); + + auto ensureFuncDecl = [&](StringRef name) { + if (!moduleOp.lookupSymbol(name)) { + auto funcType = FunctionType::get(moduleOp.getContext(), {}, {}); + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(moduleOp.getBody()); + builder.create(moduleOp.getLoc(), name, funcType) + .setPrivate(); + } + }; + + ensureFuncDecl("trace_loop_iter_begin"); + ensureFuncDecl("trace_loop_iter_end"); + + moduleOp.walk([&](Operation *op) { + if (auto whileOp = dyn_cast(op)) { + Region &afterRegion = whileOp.getAfter(); + if (afterRegion.empty()) + return; + + Block &afterBlock = afterRegion.front(); + OpBuilder builder(op->getContext()); + builder.setInsertionPointToStart(&afterBlock); + builder.create(whileOp.getLoc(), "trace_loop_iter_begin", + TypeRange(), ValueRange()); + + builder.setInsertionPoint(afterBlock.getTerminator()); + builder.create(whileOp.getLoc(), "trace_loop_iter_end", + TypeRange(), ValueRange()); + + } else if (isa(op)) { + Region &bodyRegion = op->getRegion(0); + if (bodyRegion.empty()) + return; + + Block &entryBlock = bodyRegion.front(); + OpBuilder builder(op->getContext()); + builder.setInsertionPointToStart(&entryBlock); + builder.create(op->getLoc(), "trace_loop_iter_begin", + TypeRange(), ValueRange()); + + builder.setInsertionPoint(entryBlock.getTerminator()); + builder.create(op->getLoc(), "trace_loop_iter_end", + TypeRange(), ValueRange()); + } + }); + } +}; +} // namespace + +MLIR_DECLARE_EXPLICIT_TYPE_ID(TraceLoopIterationsPass) +MLIR_DEFINE_EXPLICIT_TYPE_ID(TraceLoopIterationsPass) + +mlir::PassPluginLibraryInfo getFunctionCallCounterPassPluginInfo() { + return {MLIR_PLUGIN_API_VERSION, "TraceLoopIterationsPass", "1.0", + []() { mlir::PassRegistration(); }}; +} + +extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo +mlirGetPassPluginInfo() { + return getFunctionCallCounterPassPluginInfo(); +} diff --git a/mlir/test/compiler-course/sozonov_i_trace_loop_iterations/test.mlir b/mlir/test/compiler-course/sozonov_i_trace_loop_iterations/test.mlir new file mode 100644 index 0000000000000..0c5588099e0be --- /dev/null +++ b/mlir/test/compiler-course/sozonov_i_trace_loop_iterations/test.mlir @@ -0,0 +1,101 @@ +// RUN: mlir-opt -load-pass-plugin=%mlir_lib_dir/TraceLoopIterationsPass_Sozonov_Ilya_FIIT3_MLIR%shlibext \ +// RUN: --pass-pipeline="builtin.module(TraceLoopIterationsPass_Sozonov_Ilya_FIIT3_MLIR)" %s | FileCheck %s + +// CHECK-LABEL: func @test_scf_for +func.func @test_scf_for() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + + // CHECK: scf.for + // CHECK-NEXT: func.call @trace_loop_iter_begin + // CHECK-NEXT: %[[VAL:.*]] = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end + + scf.for %i = %c0 to %c10 step %c1 { + %0 = arith.addi %c1, %c2 : index + } + return +} + +// CHECK-LABEL: func @test_scf_while +func.func @test_scf_while() { + %c10 = arith.constant 10 : i32 + %c1 = arith.constant 1 : i32 + %init = arith.constant 0 : i32 + + // CHECK: scf.while + // CHECK-NEXT: %[[COND:.*]] = arith.cmpi + // CHECK-NEXT: scf.condition(%[[COND]]) %{{.*}} : i32 + // CHECK-NEXT: } do { + // CHECK-NEXT: ^bb0 + // CHECK-NEXT: func.call @trace_loop_iter_begin() + // CHECK-NEXT: %[[INC:.*]] = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end() + // CHECK-NEXT: scf.yield %[[INC]] : i32 + + scf.while (%x = %init) : (i32) -> (i32) { + %cond = arith.cmpi slt, %x, %c10 : i32 + scf.condition(%cond) %x : i32 + } do { + ^bb0(%x: i32): + %inc = arith.addi %x, %c1 : i32 + scf.yield %inc : i32 + } + return +} + +// CHECK-LABEL: func @test_affine_for +func.func @test_affine_for() { + %c1 = arith.constant 1 : i32 + %c2 = arith.constant 2 : i32 + + // CHECK: affine.for + // CHECK-NEXT: func.call @trace_loop_iter_begin + // CHECK-NEXT: %[[V:.*]] = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end + + affine.for %i = 0 to 10 { + %v = arith.addi %c1, %c2 : i32 + } + return +} + +// CHECK-LABEL: func @test_nested_loops +func.func @test_nested_loops() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %ci1 = arith.constant 1 : i32 + %ci2 = arith.constant 2 : i32 + + // CHECK: scf.for + // CHECK-NEXT: func.call @trace_loop_iter_begin + // CHECK-NEXT: affine.for + // CHECK-NEXT: func.call @trace_loop_iter_begin + // CHECK-NEXT: %[[V:.*]] = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end + // CHECK-NEXT: } + // CHECK-NEXT: func.call @trace_loop_iter_end + + scf.for %i = %c0 to %c10 step %c1 { + affine.for %j = 0 to 5 { + %v = arith.addi %ci1, %ci2 : i32 + } + } + return +} + +// CHECK-LABEL: func @test_no_loops +func.func @test_no_loops() { + + // CHECK-NOT: func.call @trace_loop_iter_begin + // CHECK-NOT: func.call @trace_loop_iter_end + + %c1 = arith.constant 1 : i32 + %c2 = arith.constant 2 : i32 + %v = arith.addi %c1, %c2 : i32 + return +}