diff --git a/mlir/compiler-course/drozhdinov_d/CMakeLists.txt b/mlir/compiler-course/drozhdinov_d/CMakeLists.txt new file mode 100644 index 0000000000000..e91a574a4f97c --- /dev/null +++ b/mlir/compiler-course/drozhdinov_d/CMakeLists.txt @@ -0,0 +1,16 @@ +set(Title "LoopTracePass") +set(Student "DrozhdinovD") +set(Group "FIIT1") +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/drozhdinov_d/LoopTracePass.cpp b/mlir/compiler-course/drozhdinov_d/LoopTracePass.cpp new file mode 100644 index 0000000000000..97b1a59fadd18 --- /dev/null +++ b/mlir/compiler-course/drozhdinov_d/LoopTracePass.cpp @@ -0,0 +1,84 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Tools/Plugins/PassPlugin.h" + +namespace { +class LoopTracePass + : public mlir::PassWrapper> { +public: + mlir::StringRef getArgument() const final override { + return "LoopTracePass_DrozhdinovD_FIIT1_MLIR"; + } + + mlir::StringRef getDescription() const final override { + return "Insert @trace_loop_iter_begin and @trace_loop_iter_end function " + "calls in the start and the end of loops"; + } + + void runOnOperation() override { + mlir::ModuleOp module = getOperation(); + mlir::OpBuilder builder(module.getContext()); + + Declarator("trace_loop_iter_begin", module, builder); + Declarator("trace_loop_iter_end", module, builder); + + module.walk([&](mlir::Operation *op) { + if (!isRecognizedLoop(op)) + return; + + for (mlir::Region ®ion : op->getRegions()) { + if (region.empty()) + continue; + + mlir::Block &block = region.front(); + + builder.setInsertionPointToStart(&block); + builder.create( + op->getLoc(), "trace_loop_iter_begin", mlir::TypeRange{}, + mlir::ValueRange{}); + + builder.setInsertionPoint(block.getTerminator()); + builder.create(op->getLoc(), "trace_loop_iter_end", + mlir::TypeRange{}, + mlir::ValueRange{}); + } + }); + } + +private: + bool isRecognizedLoop(mlir::Operation *op) { + return llvm::isa(op); + } + + void Declarator(mlir::StringRef name, mlir::ModuleOp module, + mlir::OpBuilder &builder) { + if (module.lookupSymbol(name)) { + return; + } + builder.setInsertionPointToStart(module.getBody()); + auto func = builder.create( + module.getLoc(), name, builder.getFunctionType({}, {})); + func.setSymVisibility("private"); + } +}; +} // namespace + +MLIR_DECLARE_EXPLICIT_TYPE_ID(LoopTracePass) +MLIR_DEFINE_EXPLICIT_TYPE_ID(LoopTracePass) + +mlir::PassPluginLibraryInfo getTraceLoopPassPluginInfo() { + return {MLIR_PLUGIN_API_VERSION, "LoopTracePass", "1.0", + []() { mlir::PassRegistration(); }}; +} + +extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo +mlirGetPassPluginInfo() { + return getTraceLoopPassPluginInfo(); +} diff --git a/mlir/test/compiler-course/drozhdinov_d/test.mlir b/mlir/test/compiler-course/drozhdinov_d/test.mlir new file mode 100644 index 0000000000000..b0a70cb864fad --- /dev/null +++ b/mlir/test/compiler-course/drozhdinov_d/test.mlir @@ -0,0 +1,78 @@ +// RUN: mlir-opt -load-pass-plugin=%mlir_lib_dir/LoopTracePass_DrozhdinovD_FIIT1_MLIR%shlibext \ +// RUN: --pass-pipeline="builtin.module(LoopTracePass_DrozhdinovD_FIIT1_MLIR)" %s | FileCheck %s + +module { + // CHECK: func.func private @trace_loop_iter_end() + // CHECK-NEXT: func.func private @trace_loop_iter_begin() + + // CHECK-LABEL: func.func @affine_for() + // CHECK: affine.for + // CHECK-NEXT: func.call @trace_loop_iter_begin() : () -> () + // CHECK-NEXT: %1 = arith.index_cast + // CHECK-NEXT: %2 = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end() : () -> () + func.func @affine_for() -> i32 { + %sum_init = arith.constant 0 : i32 + %result = affine.for %i = 0 to 10 iter_args(%sum_iter = %sum_init) -> i32 { + %i_i32 = arith.index_cast %i : index to i32 + %new_sum = arith.addi %sum_iter, %i_i32 : i32 + affine.yield %new_sum : i32 + } + return %result : i32 + } + + // CHECK-LABEL: func.func @scf_while( + // CHECK: func.call @trace_loop_iter_begin() : () -> () + // CHECK-NEXT: %c10 = arith.constant 10 : index + // CHECK-NEXT: %1 = arith.cmpi slt, %arg1, %c10 : index + // CHECK-NEXT: func.call @trace_loop_iter_end() : () -> () + // CHECK-NEXT: scf.condition(%1) %arg1 : index + // CHECK-NEXT: } do { + // CHECK-NEXT: ^bb0(%arg1: index): + // CHECK-NEXT: func.call @trace_loop_iter_begin() : () -> () + // CHECK-NEXT: %1 = arith.addi %arg1, %c1 : index + // CHECK-NEXT: func.call @trace_loop_iter_end() : () -> () + func.func @scf_while(%arg0: index) { + %one = arith.constant 1 : index + %res = scf.while (%i = %arg0) : (index) -> (index) { + %limit = arith.constant 10 : index + %cond = arith.cmpi slt, %i, %limit : index + scf.condition(%cond) %i : index + } do { + ^bb0(%i_curr: index): + %inc = arith.addi %i_curr, %one : index + scf.yield %inc : index + } + return + } + + // CHECK-LABEL: func.func @scf_for() + // CHECK: scf.for + // CHECK-NEXT: func.call @trace_loop_iter_begin() : () -> () + // CHECK-NEXT: %0 = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end() : () -> () + func.func @scf_for() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + scf.for %i = %c0 to %c10 step %c1 { + %x = arith.addi %i, %c1 : index + } + return + } + + // CHECK-LABEL: func.func @scf_parallel() + // CHECK: scf.parallel + // CHECK-NEXT: func.call @trace_loop_iter_begin() : () -> () + // CHECK-NEXT: %0 = arith.addi + // CHECK-NEXT: func.call @trace_loop_iter_end() : () -> () + func.func @scf_parallel() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + scf.parallel (%i) = (%c0) to (%c10) step (%c1) { + %x = arith.addi %i, %c1 : index + } + return + } +}