Хохлов Андрей. Лабораторная работа 4: MLIR. Вариант 4

mlir/compiler-course/khokhlov_a_mlir/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(Title "LoopNestDepth")
+set(Student "Khokhlov_Andrey")
+set(Group "FIIT2")
+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/khokhlov_a_mlir/LoopNestDepthPass.cpp

@@ -0,0 +1,111 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Tools/Plugins/PassPlugin.h"
+#include "llvm/Support/raw_ostream.h"
+
+namespace {
+
+int computeRegionNestingDepth(mlir::Region &Region) {
+  int MaxDepth = 0;
+
+  // Обходим все блоки в регионе
+  for (mlir::Block &block : Region) {
+    for (mlir::Operation &op : block) {
+      // Проверяем, является ли операция циклом или условным оператором
+      if (op.hasTrait<mlir::OpTrait::IsTerminator>() ||
+          !(mlir::isa<mlir::scf::ForOp, mlir::scf::WhileOp,
+                      mlir::affine::AffineForOp, mlir::scf::IfOp,
+                      mlir::affine::AffineIfOp>(op))) {
+        continue;
+      }
+
+      // Начальная глубина для текущей операции
+      int opDepth = 1;
+
+      // Рекурсивно вычисляем максимальную глубину вложенных регионов
+      int maxSubRegionDepth = 0;
+      for (mlir::Region &subRegion : op.getRegions()) {
+        int subDepth = computeRegionNestingDepth(subRegion);
+        maxSubRegionDepth = std::max(maxSubRegionDepth, subDepth);
+      }
+
+      // Общая глубина = 1 (текущая операция) + максимальная глубина подрегионов
+      opDepth += maxSubRegionDepth;
+
+      // Обновляем максимальную глубину
+      MaxDepth = std::max(MaxDepth, opDepth);
+    }
+  }
+  return MaxDepth;
+}
+
+class LoopNestDepthPass
+    : public mlir::PassWrapper<LoopNestDepthPass,
+                               mlir::OperationPass<mlir::ModuleOp>> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LoopNestDepthPass)
+
+  mlir::StringRef getArgument() const override { return "loop-nest-depth"; }
+  mlir::StringRef getDescription() const override {
+    return "Computes the maximum nesting depth of loops and conditionals in "
+           "functions";
+  }
+
+  void runOnOperation() override {
+    mlir::ModuleOp Module = getOperation();
+    mlir::OpBuilder Builder(Module.getContext());
+
+    // Обходим все функции в модуле
+    Module.walk([&](mlir::func::FuncOp Func) {
+      llvm::SmallVector<int64_t> LoopDepths;
+
+      // Обходим все операции в функции
+      Func.walk([&](mlir::Operation *Op) {
+        if (mlir::isa<mlir::scf::ForOp, mlir::scf::WhileOp,
+                      mlir::affine::AffineForOp>(Op)) {
+          // Начальная глубина цикла = 1
+          int Depth = 1;
+
+          // Проверяем вложенные регионы
+          int MaxNestedDepth = 0;
+          for (mlir::Region &Region : Op->getRegions()) {
+            int regionDepth = computeRegionNestingDepth(Region);
+            MaxNestedDepth = std::max(MaxNestedDepth, regionDepth);
+          }
+
+          // Общая глубина = 1 + максимальная глубина вложенных регионов
+          LoopDepths.push_back(Depth + MaxNestedDepth);
+        }
+      });
+
+      // Устанавливаем атрибут, если найдены циклы
+      if (!LoopDepths.empty()) {
+        Func->setAttr("my_loop_depths", Builder.getI64ArrayAttr(LoopDepths));
+        auto &os = llvm::outs();
+        os << "Function '" << Func.getName()
+                     << "': processed. Loop depths: ";
+        llvm::interleaveComma(LoopDepths, os);
+        os << "\n";
+      } else {
+        os << "Function '" << Func.getName() << "': no loops found\n";
+      }
+    });
+  }
+};
+
+} // namespace
+
+static mlir::PassPluginLibraryInfo getPassPluginInfo() {
+  return {MLIR_PLUGIN_API_VERSION, "LoopNestDepthPass", "1.0.0",
+          []() { mlir::PassRegistration<LoopNestDepthPass>(); }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo
+mlirGetPassPluginInfo() {
+  return getPassPluginInfo();
+}

mlir/test/compiler-course/khokhlov_a_mlir/test.mlir

@@ -0,0 +1,194 @@
+// RUN: mlir-opt -load-pass-plugin=%mlir_lib_dir/LoopNestDepth_Khokhlov_Andrey_FIIT2_MLIR%shlibext \
+// RUN: --pass-pipeline="builtin.module(loop-nest-depth)" %s | FileCheck %s
+
+// CHECK: module {
+
+// CHECK-LABEL:  func.func @no_loops() {
+// CHECK-NOT:      my_loop_depths
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @no_loops() {
+  return
+}
+
+// CHECK-LABEL:  func.func @simple_scf_for() attributes {my_loop_depths = [1]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @simple_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 {
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @simple_affine_for() attributes {my_loop_depths = [1]} {
+// CHECK-NEXT:     affine.for %arg0 = 0 to 10 {
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @simple_affine_for() {
+  affine.for %i = 0 to 10 {
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @scf_for_with_if() attributes {my_loop_depths = [2]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     %true = arith.constant true
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:       scf.if %true {
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @scf_for_with_if() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  %cond = arith.constant true
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.if %cond {
+    }
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @nested_scf_for() attributes {my_loop_depths = [1, 2]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c5 = arith.constant 5 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:       scf.for %arg1 = %c0 to %c5 step %c1 {
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @nested_scf_for() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c5 = arith.constant 5 : index
+  %c1 = arith.constant 1 : index
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.for %j = %c0 to %c5 step %c1 {
+    }
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @deeply_nested_if_in_for() attributes {my_loop_depths = [4]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     %true = arith.constant true
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:       scf.if %true {
+// CHECK-NEXT:         scf.if %true {
+// CHECK-NEXT:           scf.if %true {
+// CHECK-NEXT:           }
+// CHECK-NEXT:         }
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @deeply_nested_if_in_for() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  %cond = arith.constant true
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.if %cond {
+      scf.if %cond {
+        scf.if %cond {
+        }
+      }
+    }
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @two_independent_loops() attributes {my_loop_depths = [2, 1]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     %true = arith.constant true
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:       scf.if %true {
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:     affine.for %arg0 = 0 to 5 {
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @two_independent_loops() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  %cond = arith.constant true
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.if %cond {
+    }
+  }
+  affine.for %j = 0 to 5 {
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @loop_if_loop() attributes {my_loop_depths = [1, 3]} {
+// CHECK-NEXT:     %c0 = arith.constant 0 : index
+// CHECK-NEXT:     %c10 = arith.constant 10 : index
+// CHECK-NEXT:     %c1 = arith.constant 1 : index
+// CHECK-NEXT:     %true = arith.constant true
+// CHECK-NEXT:     scf.for %arg0 = %c0 to %c10 step %c1 {
+// CHECK-NEXT:       scf.if %true {
+// CHECK-NEXT:         affine.for %arg1 = 0 to 5 {
+// CHECK-NEXT:         }
+// CHECK-NEXT:       }
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+func.func @loop_if_loop() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  %cond = arith.constant true
+  scf.for %i = %c0 to %c10 step %c1 {
+    scf.if %cond {
+      affine.for %j = 0 to 5 {
+      }
+    }
+  }
+  return
+}
+
+// CHECK-LABEL:  func.func @another_func_no_loops() {
+// CHECK-NOT:      my_loop_depths
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+
+// CHECK-LABEL: func.func @another_func_with_loop() attributes {my_loop_depths = [1]} {
+// CHECK-NOT:      my_loop_depths 
+// CHECK-NEXT:     affine.for %arg0 = 0 to 1 {
+// CHECK-NEXT:     }
+// CHECK-NEXT:     return
+// CHECK-NEXT:   }
+// CHECK-NEXT: }
+module {
+  func.func @another_func_no_loops() {
+    return
+  }
+  func.func @another_func_with_loop() {
+    affine.for %i = 0 to 1 {
+    }
+    return
+  }
+}