Колодкин Григорий. Лабораторная работа №4: MLIR. Вариант 5.

mlir/compiler-course/kologkin_g_lab_4_var_5/CMakeLists.txt

@@ -0,0 +1,16 @@
+set(Title "ReplaceCeilWithFloorNegPlugin")
+set(Student "KolodkinGrigorii")
+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/kologkin_g_lab_4_var_5/ReplaceCeilWithFloorNegPlugin.cpp

@@ -0,0 +1,75 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Tools/Plugins/PassPlugin.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+
+void replaceCeilWithNegFloor(ModuleOp module) {
+
+  for (auto func : module.getOps<mlir::func::FuncOp>()) {
+
+    func.walk([&](Operation *op) {
+      if (auto ceilOp = dyn_cast<math::CeilOp>(op)) {
+        Location loc = ceilOp->getLoc();
+        Value input = ceilOp->getOperand(0);
+        auto type = input.getType();
+
+        if (auto vectorType = type.dyn_cast<VectorType>()) {
+          if (!vectorType.getElementType().isa<FloatType>())
+            return;
+        } else if (!type.isa<FloatType>()) {
+          return;
+        }
+
+        OpBuilder builder(ceilOp);
+
+        auto neg1 = builder.create<arith::NegFOp>(loc, type, input);
+
+        auto floorNeg = builder.create<math::FloorOp>(loc, type, neg1);
+
+        auto neg2 = builder.create<arith::NegFOp>(loc, type, floorNeg);
+
+        ceilOp->replaceAllUsesWith(neg2);
+
+        ceilOp->erase();
+      }
+    });
+  }
+}
+
+class ReplaceCeilWithFloorNegPlugin
+    : public PassWrapper<ReplaceCeilWithFloorNegPlugin,
+                         OperationPass<ModuleOp>> {
+public:
+  StringRef getArgument() const final {
+    return "ReplaceCeilWithFloorNegPlugin_KolodkinGrigorii_FIIT3_MLIR";
+  }
+  StringRef getDescription() const final { return "Description pass"; }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    replaceCeilWithNegFloor(module);
+  }
+};
+} // namespace
+
+MLIR_DECLARE_EXPLICIT_TYPE_ID(ReplaceCeilWithFloorNegPlugin)
+MLIR_DEFINE_EXPLICIT_TYPE_ID(ReplaceCeilWithFloorNegPlugin)
+
+mlir::PassPluginLibraryInfo getFunctionCallCounterPassPluginInfo() {
+  return {MLIR_PLUGIN_API_VERSION, "ReplaceCeilWithFloorNegPlugin", "1.0",
+          []() { PassRegistration<ReplaceCeilWithFloorNegPlugin>(); }};
+}
+
+extern "C" LLVM_ATTRIBUTE_WEAK mlir::PassPluginLibraryInfo
+mlirGetPassPluginInfo() {
+  return getFunctionCallCounterPassPluginInfo();
+}

mlir/test/compiler-course/kologkin_g_lab_4_var_5/test.mlir

@@ -0,0 +1,62 @@
+// RUN: mlir-opt -load-pass-plugin=%mlir_lib_dir/ReplaceCeilWithFloorNegPlugin_KolodkinGrigorii_FIIT3_MLIR%shlibext \
+// RUN: --pass-pipeline="builtin.module(ReplaceCeilWithFloorNegPlugin_KolodkinGrigorii_FIIT3_MLIR)" %s | FileCheck %s
+
+//CHECK: func.func @test_simple(%arg0: f32) -> f32 {
+//CHECK-NEXT: %0 = arith.negf %arg0 : f32
+//CHECK-NEXT: %1 = math.floor %0 : f32
+//CHECK-NEXT: %2 = arith.negf %1 : f32
+//CHECK-NEXT: return %2 : f32
+//CHECK-NEXT: }
+
+func.func @test_simple(%arg0: f32) -> f32 {
+  %0 = math.ceil %arg0 : f32
+  return %0 : f32
+}
+
+//CHECK: func.func @test_single_arg(%arg0: f32) -> f32 {
+//CHECK-NEXT: %0 = arith.negf %arg0 : f32
+//CHECK-NEXT: %1 = math.floor %0 : f32
+//CHECK-NEXT: %2 = arith.negf %1 : f32
+//CHECK-NEXT: %3 = arith.negf %2 : f32
+//CHECK-NEXT: %4 = math.floor %3 : f32
+//CHECK-NEXT: %5 = arith.negf %4 : f32
+//CHECK-NEXT: return %5 : f32
+//CHECK-NEXT: }
+
+func.func @test_single_arg(%arg: f32) -> f32 {
+  %c1 = math.ceil %arg : f32
+  %c2 = math.ceil %c1 : f32
+  return %c2 : f32
+}
+
+//CHECK: func.func @test_vector_after_ceil(%arg0: vector<2xf32>) -> vector<2xf32> {
+//CHECK-NEXT: %0 = arith.negf %arg0 : vector<2xf32>
+//CHECK-NEXT: %1 = math.floor %0 : vector<2xf32>
+//CHECK-NEXT: %2 = arith.negf %1 : vector<2xf32>
+//CHECK-NEXT: %3 = arith.mulf %2, %arg0 : vector<2xf32>
+//CHECK-NEXT: return %3 : vector<2xf32>
+//CHECK-NEXT: }
+
+func.func @test_vector_after_ceil(%vec: vector<2xf32>) -> vector<2xf32> {
+  %ceil_vec = math.ceil %vec : vector<2xf32>
+  %result = arith.mulf %ceil_vec, %vec : vector<2xf32>
+  return %result : vector<2xf32>
+}
+
+//CHECK: func.func @test_two_args(%arg0: f64, %arg1: f64) -> f64 {
+//CHECK-NEXT: %0 = arith.negf %arg0 : f64
+//CHECK-NEXT: %1 = math.floor %0 : f64
+//CHECK-NEXT: %2 = arith.negf %1 : f64
+//CHECK-NEXT: %3 = arith.negf %arg1 : f64
+//CHECK-NEXT: %4 = math.floor %3 : f64
+//CHECK-NEXT: %5 = arith.negf %4 : f64
+//CHECK-NEXT: %6 = arith.mulf %2, %5 : f64
+//CHECK-NEXT: return %6 : f64
+//CHECK-NEXT: }
+
+func.func @test_two_args(%x: f64, %y: f64) -> f64 {
+  %ceil_x = math.ceil %x : f64
+  %ceil_y = math.ceil %y : f64
+  %sum = arith.mulf %ceil_x, %ceil_y : f64
+  return %sum : f64
+}