Созонов Илья. Лабораторная работа 4. MLIR. Вариант 1.

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)

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<TraceLoopIterationsPass, OperationPass<ModuleOp>> {
+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<func::FuncOp>(name)) {
+        auto funcType = FunctionType::get(moduleOp.getContext(), {}, {});
+        OpBuilder::InsertionGuard guard(builder);
+        builder.setInsertionPointToStart(moduleOp.getBody());
+        builder.create<func::FuncOp>(moduleOp.getLoc(), name, funcType)
+            .setPrivate();
+      }
+    };
+
+    ensureFuncDecl("trace_loop_iter_begin");
+    ensureFuncDecl("trace_loop_iter_end");
+
+    moduleOp.walk([&](Operation *op) {
+      if (auto whileOp = dyn_cast<scf::WhileOp>(op)) {
+        Region &afterRegion = whileOp.getAfter();
+        if (afterRegion.empty())
+          return;
+
+        Block &afterBlock = afterRegion.front();
+        OpBuilder builder(op->getContext());
+        builder.setInsertionPointToStart(&afterBlock);
+        builder.create<func::CallOp>(whileOp.getLoc(), "trace_loop_iter_begin",
+                                     TypeRange(), ValueRange());
+
+        builder.setInsertionPoint(afterBlock.getTerminator());
+        builder.create<func::CallOp>(whileOp.getLoc(), "trace_loop_iter_end",
+                                     TypeRange(), ValueRange());
+
+      } else if (isa<scf::ForOp, affine::AffineForOp>(op)) {
+        Region &bodyRegion = op->getRegion(0);
+        if (bodyRegion.empty())
+          return;
+
+        Block &entryBlock = bodyRegion.front();
+        OpBuilder builder(op->getContext());
+        builder.setInsertionPointToStart(&entryBlock);
+        builder.create<func::CallOp>(op->getLoc(), "trace_loop_iter_begin",
+                                     TypeRange(), ValueRange());
+
+        builder.setInsertionPoint(entryBlock.getTerminator());
+        builder.create<func::CallOp>(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<TraceLoopIterationsPass>(); }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo
+mlirGetPassPluginInfo() {
+  return getFunctionCallCounterPassPluginInfo();
+}

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