[ReuseIR] add proj lowering (#18)

SchrodingerZhu · Aug 21, 2024 · 3a68cd0 · 3a68cd0
1 parent e9d242a
commit 3a68cd0
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 3 deletions.
diff --git a/reuse-mlir/src/cxx/IR/ReuseIROps.cpp b/reuse-mlir/src/cxx/IR/ReuseIROps.cpp
@@ -44,6 +44,9 @@ mlir::reuse_ir::LogicalResult ValueToRefOp::verify() {
 // ProjOp
 mlir::reuse_ir::LogicalResult ProjOp::verify() {
   RefType input = getObject().getType();
+  if (input.getFreezingKind() != getType().getFreezingKind())
+    return emitOpError(
+        "must return a reference with the same freezing kind as the input");
   if (!isProjectable(input.getPointee()))
     return emitOpError(
         "must operate on a reference to a composite or an array");
@@ -57,7 +60,7 @@ mlir::reuse_ir::LogicalResult ProjOp::verify() {
           .Case<ArrayType>([&](const ArrayType &ty) {
             if (getIndex() >= ty.getSizes()[0])
               return mlir::Type{};
-            if (ty.getSizes().size() == 0)
+            if (ty.getSizes().size() == 1)
               return ty.getElementType();
             return cast<mlir::Type>(ArrayType::get(
                 getContext(), ty.getElementType(), ty.getSizes().drop_front()));

diff --git a/reuse-mlir/src/cxx/Passes/ConvertReuseIRToLLVM.cpp b/reuse-mlir/src/cxx/Passes/ConvertReuseIRToLLVM.cpp
@@ -39,6 +39,44 @@ class ReuseIRConvPatternWithLayoutCache : public mlir::OpConversionPattern<Op> {
         cache(cache) {}
 };
 
+class ProjOpLowering : public ReuseIRConvPatternWithLayoutCache<ProjOp> {
+public:
+  using ReuseIRConvPatternWithLayoutCache::ReuseIRConvPatternWithLayoutCache;
+
+  mlir::reuse_ir::LogicalResult matchAndRewrite(
+      ProjOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter &rewriter) const override final {
+    auto ptrTy = LLVM::LLVMPointerType::get(getContext());
+    if (auto pointee =
+            dyn_cast<CompositeType>(op.getObject().getType().getPointee())) {
+      const CompositeLayout &layout = cache.get(pointee);
+      rewriter.replaceOpWithNewOp<LLVM::GEPOp>(
+          op, ptrTy, layout.getLLVMType(getLLVMTypeConverter()),
+          adaptor.getObject(),
+          llvm::ArrayRef<LLVM::GEPArg>{0,
+                                       layout.getField(op.getIndex()).index});
+      return LogicalResult::success();
+    }
+    if (auto pointee =
+            dyn_cast<ArrayType>(op.getObject().getType().getPointee())) {
+      mlir::Type innerTy;
+      if (pointee.getSizes().size() == 0)
+        innerTy = pointee.getElementType();
+      else
+        innerTy = ArrayType::get(getContext(), pointee.getElementType(),
+                                 pointee.getSizes().drop_front());
+      auto ty = typeConverter->convertType(innerTy);
+      if (!ty)
+        return LogicalResult::failure();
+      rewriter.replaceOpWithNewOp<LLVM::GEPOp>(
+          op, ptrTy, ty, adaptor.getObject(),
+          llvm::ArrayRef<LLVM::GEPArg>{op.getIndex()});
+      return LogicalResult::success();
+    }
+    return LogicalResult::failure();
+  }
+};
+
 class ValueToRefOpLowering
     : public ReuseIRConvPatternWithLayoutCache<ValueToRefOp> {
 public:
@@ -228,8 +266,8 @@ void ConvertReuseIRToLLVMPass::runOnOperation() {
   mlir::populateFuncToLLVMConversionPatterns(converter, patterns);
   patterns.add<IncOpLowering, AllocOpLowering, FreeOpLowering>(converter,
                                                                &getContext());
-  patterns.add<BorrowOpLowering, ValueToRefOpLowering>(cache, converter,
-                                                       &getContext());
+  patterns.add<BorrowOpLowering, ValueToRefOpLowering, ProjOpLowering>(
+      cache, converter, &getContext());
   mlir::ConversionTarget target(getContext());
   target.addLegalDialect<mlir::LLVM::LLVMDialect>();
   target.addLegalOp<mlir::ModuleOp>();

diff --git a/reuse-mlir/test/integration/basic/invalid_projection_kind.mlir b/reuse-mlir/test/integration/basic/invalid_projection_kind.mlir
@@ -0,0 +1,13 @@
+// RUN: %not %reuse-opt %s 2>&1 | %FileCheck %s
+!test = !reuse_ir.composite<!reuse_ir.composite<i32, i32, f128>, i32>
+module @test {
+    func.func @projection(%0: !reuse_ir.rc<!test, nonatomic, nonfreezing>) {
+        %1 = reuse_ir.borrow %0 : 
+            !reuse_ir.rc<!test, nonatomic, nonfreezing>
+            -> !reuse_ir.ref<!test, nonfreezing>
+        // CHECK: error: 'reuse_ir.proj' op must return a reference with the same freezing kind as the input
+        %2 = reuse_ir.proj %1[0] : 
+            !reuse_ir.ref<!test, nonfreezing> -> !reuse_ir.ref<!reuse_ir.composite<i32, i32, f128>, frozen>
+        return
+    }
+}
diff --git a/reuse-mlir/test/integration/lowering/proj_op_lowering.mlir b/reuse-mlir/test/integration/lowering/proj_op_lowering.mlir
@@ -0,0 +1,33 @@
+// RUN: %reuse-opt %s -convert-reuse-ir-to-llvm | %FileCheck %s
+!test = !reuse_ir.composite<!reuse_ir.composite<i32, i32, f128>, i32>
+!array = !reuse_ir.array<!test, 16, 16>
+module @test attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<f80, dense<128> : vector<2xi64>>, #dlti.dl_entry<i128, dense<128> : vector<2xi64>>, #dlti.dl_entry<i32, dense<32> : vector<2xi64>>, #dlti.dl_entry<f128, dense<128> : vector<2xi64>>, #dlti.dl_entry<f64, dense<64> : vector<2xi64>>, #dlti.dl_entry<f16, dense<16> : vector<2xi64>>, #dlti.dl_entry<i1, dense<8> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<270>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i8, dense<8> : vector<2xi64>>, #dlti.dl_entry<i16, dense<16> : vector<2xi64>>, #dlti.dl_entry<!llvm.ptr<272>, dense<64> : vector<4xi64>>, #dlti.dl_entry<!llvm.ptr<271>, dense<32> : vector<4xi64>>, #dlti.dl_entry<i64, dense<64> : vector<2xi64>>, #dlti.dl_entry<"dlti.stack_alignment", 128 : i64>, #dlti.dl_entry<"dlti.endianness", "little">>, llvm.data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"} 
+{
+    func.func @composite_projection(%0: !reuse_ir.rc<!test, nonatomic, nonfreezing>) {
+        %1 = reuse_ir.borrow %0 : 
+            !reuse_ir.rc<!test, nonatomic, nonfreezing>
+            -> !reuse_ir.ref<!test, nonfreezing>
+        // CHECK: %[[REG0:[0-9a-z]+]] = llvm.getelementptr %{{[0-9a-z]+}}[0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>
+        %2 = reuse_ir.proj %1[0] : 
+            !reuse_ir.ref<!test, nonfreezing> -> !reuse_ir.ref<!reuse_ir.composite<i32, i32, f128>, nonfreezing>
+        // CHECK: %{{[0-9a-z]+}} = llvm.getelementptr %[[REG0]][0, 3] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(i32, i32, array<8 x i8>, f128)>
+        %3 = reuse_ir.proj %2[2] : 
+            !reuse_ir.ref<!reuse_ir.composite<i32, i32, f128>, nonfreezing> -> !reuse_ir.ref<f128, nonfreezing>
+        return
+    }
+    func.func @array_projection(%0: !reuse_ir.rc<!array, nonatomic, nonfreezing>) {
+        %1 = reuse_ir.borrow %0 : 
+            !reuse_ir.rc<!array, nonatomic, nonfreezing>
+            -> !reuse_ir.ref<!array, nonfreezing>
+        // CHECK: %[[REG0:[0-9a-z]+]] = llvm.getelementptr %{{[0-9a-z]+}}[13] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<16 x struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>>
+        %2 = reuse_ir.proj %1[13] : 
+            !reuse_ir.ref<!array, nonfreezing> -> !reuse_ir.ref<!reuse_ir.array<!test, 16>, nonfreezing>
+        // CHECK: %[[REG1:[0-9a-z]+]] = llvm.getelementptr %[[REG0]][7] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>
+        %3 = reuse_ir.proj %2[7] : 
+            !reuse_ir.ref<!reuse_ir.array<!test, 16>, nonfreezing> -> !reuse_ir.ref<!test, nonfreezing>
+        // CHECK: %{{[0-9a-z]+}} = llvm.getelementptr %[[REG1]][0, 1] : (!llvm.ptr) -> !llvm.ptr, !llvm.struct<(struct<(i32, i32, array<8 x i8>, f128)>, i32, array<12 x i8>)>
+        %4 = reuse_ir.proj %3[1] : 
+            !reuse_ir.ref<!test, nonfreezing>-> !reuse_ir.ref<i32, nonfreezing>
+        return
+    }
+}