Дрождинов Дмитрий. Лабораторная работа 4: MLIR. Вариант 1.

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)

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<LoopTracePass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+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 &region : op->getRegions()) {
+        if (region.empty())
+          continue;
+
+        mlir::Block &block = region.front();
+
+        builder.setInsertionPointToStart(&block);
+        builder.create<mlir::func::CallOp>(
+            op->getLoc(), "trace_loop_iter_begin", mlir::TypeRange{},
+            mlir::ValueRange{});
+
+        builder.setInsertionPoint(block.getTerminator());
+        builder.create<mlir::func::CallOp>(op->getLoc(), "trace_loop_iter_end",
+                                           mlir::TypeRange{},
+                                           mlir::ValueRange{});
+      }
+    });
+  }
+
+private:
+  bool isRecognizedLoop(mlir::Operation *op) {
+    return llvm::isa<mlir::scf::ForOp, mlir::scf::WhileOp,
+                     mlir::scf::ParallelOp, mlir::affine::AffineForOp>(op);
+  }
+
+  void Declarator(mlir::StringRef name, mlir::ModuleOp module,
+                  mlir::OpBuilder &builder) {
+    if (module.lookupSymbol<mlir::func::FuncOp>(name)) {
+      return;
+    }
+    builder.setInsertionPointToStart(module.getBody());
+    auto func = builder.create<mlir::func::FuncOp>(
+        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<LoopTracePass>(); }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo
+mlirGetPassPluginInfo() {
+  return getTraceLoopPassPluginInfo();
+}

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
+  }
+}