diff --git a/mlir/compiler-course/rezantseva_mlir/CMakeLists.txt b/mlir/compiler-course/rezantseva_mlir/CMakeLists.txt new file mode 100644 index 0000000000000..8f3208f137aad --- /dev/null +++ b/mlir/compiler-course/rezantseva_mlir/CMakeLists.txt @@ -0,0 +1,16 @@ +set(Title "MaxDepthPass") +set(Student "RezantsevaAnastasia") +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/rezantseva_mlir/rezantseva_mlir.cpp b/mlir/compiler-course/rezantseva_mlir/rezantseva_mlir.cpp new file mode 100644 index 0000000000000..c244a42e336bb --- /dev/null +++ b/mlir/compiler-course/rezantseva_mlir/rezantseva_mlir.cpp @@ -0,0 +1,112 @@ +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Tools/Plugins/PassPlugin.h" +#include +#include + +using namespace mlir; + +namespace { + +class MaxDepthPass + : public PassWrapper> { +private: + bool isLoopOp(Operation *op) { + return isa(op); + } + + bool isNestingOp(Operation *op) { + return isa(op); + } + + unsigned computeDepth(Operation *op, unsigned currentDepth = 0) { + unsigned maxDepth = currentDepth; + + if (isNestingOp(op)) { + maxDepth = std::max(maxDepth, currentDepth + 1); + } + + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + maxDepth = std::max( + maxDepth, computeDepth(&nestedOp, + currentDepth + (isNestingOp(op) ? 1 : 0))); + } + } + } + + return maxDepth; + } + +public: + StringRef getArgument() const final { + return "MaxDepthPass_RezantsevaAnastasia_FIIT1_MLIR"; + } + + StringRef getDescription() const final { + return "The pass counts the max depth of control flow operations nests in each " + "loop."; + } + + void runOnOperation() override { + func::FuncOp func = getOperation(); + std::vector loopDepths; + + func.walk([&](Operation *op) { + if (isLoopOp(op)) { + Operation *parent = op->getParentOp(); + bool isNestedInLoop = false; + while (parent && !isa(parent)) { + if (isLoopOp(parent)) { + isNestedInLoop = true; + break; + } + parent = parent->getParentOp(); + } + if (!isNestedInLoop) { + unsigned depth = computeDepth(op); + if (depth > 0) { + loopDepths.push_back(depth); + } + } + } + }); + + if (loopDepths.empty()) { + func->setAttr("my_loop_depths", ArrayAttr::get(func.getContext(), {})); + } else { + std::sort(loopDepths.begin(), loopDepths.end(), std::greater()); + SmallVector depthAttrs; + for (unsigned depth : loopDepths) { + depthAttrs.push_back( + IntegerAttr::get(IntegerType::get(func.getContext(), 64), depth)); + } + func->setAttr("my_loop_depths", + ArrayAttr::get(func.getContext(), depthAttrs)); + } + } +}; + +} // namespace + +MLIR_DECLARE_EXPLICIT_TYPE_ID(MaxDepthPass) +MLIR_DEFINE_EXPLICIT_TYPE_ID(MaxDepthPass) + +mlir::PassPluginLibraryInfo getMaxDepthPassPluginInfo() { + return {MLIR_PLUGIN_API_VERSION, "MaxDepthPass", "1.0", + []() { mlir::PassRegistration(); }}; +} + +extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo +mlirGetPassPluginInfo() { + return getMaxDepthPassPluginInfo(); +} diff --git a/mlir/test/compiler-course/rezantseva_mlir/rezantseva_test.mlir b/mlir/test/compiler-course/rezantseva_mlir/rezantseva_test.mlir new file mode 100644 index 0000000000000..3e5b3dc864189 --- /dev/null +++ b/mlir/test/compiler-course/rezantseva_mlir/rezantseva_test.mlir @@ -0,0 +1,178 @@ +// RUN: mlir-opt -load-pass-plugin=%mlir_lib_dir/MaxDepthPass_RezantsevaAnastasia_FIIT1_MLIR%shlibext \ +// RUN: --pass-pipeline="builtin.module(func.func(MaxDepthPass_RezantsevaAnastasia_FIIT1_MLIR))" %s | FileCheck %s +#set = affine_set<(d0): (d0 >= 0)> + +// CHECK: func.func @only_affine_if() attributes {my_loop_depths = []} +// CHECK-NEXT: %c0 = arith.constant 0 : index +// CHECK-NEXT: affine.if #set(%c0) { +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } +func.func @only_affine_if() { + %c0 = arith.constant 0 : index + affine.if #set(%c0) { + } + return +} + +// CHECK: func.func @single_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 @single_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: func.func @affine_for_in_affine_if() attributes {my_loop_depths = [1]} +// CHECK-NEXT: %c0 = arith.constant 0 : index +// CHECK-NEXT: %c10 = arith.constant 10 : index +// CHECK-NEXT: affine.if #set(%c0) { +// CHECK-NEXT: affine.for %arg0 = 0 to 10 +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } +func.func @affine_for_in_affine_if() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + affine.if #set(%c0) { + affine.for %arg0 = 0 to 10 { + } + } + return +} + +// CHECK: func.func @single_scf_while() 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: %0 = scf.while (%arg0 = %c0) : (index) -> index { +// CHECK-NEXT: %1 = arith.cmpi slt, %arg0, %c10 : index +// CHECK-NEXT: scf.condition(%1) %arg0 : index +// CHECK-NEXT: } do { +// CHECK-NEXT: ^bb0(%arg0: index): +// CHECK-NEXT: %1 = arith.addi %arg0, %c1 : index +// CHECK-NEXT: scf.yield %1 : index +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } + +func.func @single_scf_while() { + %c0 = arith.constant 0 : index + %c10 = arith.constant 10 : index + %c1 = arith.constant 1 : index + %res = scf.while (%i = %c0) : (index) -> index { + %cond = arith.cmpi slt, %i, %c10 : index + scf.condition(%cond) %i : index + } do { + ^bb0(%i: index): + %next = arith.addi %i, %c1 : index + scf.yield %next : index + } + return +} + + + +// CHECK: 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 %{{.*}} = %c0 to %c10 step %c1 { +// CHECK-NEXT: scf.if %{{.*}} { +// 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: 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: %{{.*}} = arith.constant true +// CHECK-NEXT: scf.for %{{.*}} = %c0 to %c10 step %c1 { +// CHECK-NEXT: scf.if %{{.*}} { +// CHECK-NEXT: scf.if %{{.*}} { +// CHECK-NEXT: scf.if %{{.*}} { +// 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: func.func @nested_loops() attributes {my_loop_depths = [4, 2]} { +// CHECK-NEXT: %c0 = arith.constant 0 : index +// CHECK-NEXT: %c20 = arith.constant 20 : index +// CHECK-NEXT: %c2 = arith.constant 2 : index +// CHECK-NEXT: scf.for %arg0 = %c0 to %c20 step %c2 { +// CHECK-NEXT: scf.for %arg1 = %c0 to %c20 step %c2 { +// CHECK-NEXT: scf.for %arg2 = %c0 to %c20 step %c2 { +// CHECK-NEXT: scf.for %arg3 = %c0 to %c20 step %c2 { +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: scf.for %arg0 = %c0 to %c20 step %c2 { +// CHECK-NEXT: scf.for %arg1 = %c0 to %c20 step %c2 { +// CHECK-NEXT: } +// CHECK-NEXT: } +// CHECK-NEXT: return +// CHECK-NEXT: } + +func.func @nested_loops() { + %c0 = arith.constant 0 : index + %c20 = arith.constant 20 : index + %c2 = arith.constant 2 : index + scf.for %outer = %c0 to %c20 step %c2 { + scf.for %inner1 = %c0 to %c20 step %c2 { + scf.for %inner2 = %c0 to %c20 step %c2 { + scf.for %inner3 = %c0 to %c20 step %c2 { + } + } + } + } + scf.for %outer2 = %c0 to %c20 step %c2 { + scf.for %inner4 = %c0 to %c20 step %c2 { + } + } + return +}