[MLIR] MemRefToEmitC: convert memref.load/store to emitc.subscript

Pull Request: #119
Xilinx · Feb 28, 2024 · bb4596e · bb4596e
1 parent c938fb8
commit bb4596e
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 0 deletions.
diff --git a/mlir/lib/Conversion/MemRefToEmitC/MemRefToEmitC.cpp b/mlir/lib/Conversion/MemRefToEmitC/MemRefToEmitC.cpp
@@ -54,6 +54,34 @@ bool isSignatureLegal(FunctionType ty) {
   return isLegal(llvm::concat<const Type>(ty.getInputs(), ty.getResults()));
 }
 
+struct ConvertLoad final : public OpConversionPattern<memref::LoadOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::LoadOp op, OpAdaptor operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    rewriter.replaceOpWithNewOp<emitc::SubscriptOp>(op, operands.getMemref(),
+                                                    operands.getIndices());
+    return success();
+  }
+};
+
+struct ConvertStore final : public OpConversionPattern<memref::StoreOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::StoreOp op, OpAdaptor operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    auto subscript = rewriter.create<emitc::SubscriptOp>(
+        op.getLoc(), operands.getMemref(), operands.getIndices());
+    rewriter.replaceOpWithNewOp<emitc::AssignOp>(op, subscript,
+                                                 operands.getValue());
+    return success();
+  }
+};
+
 struct ConvertMemRefToEmitCPass
     : public impl::ConvertMemRefToEmitCBase<ConvertMemRefToEmitCPass> {
   void runOnOperation() override {
@@ -91,6 +119,7 @@ struct ConvertMemRefToEmitCPass
     target.addDynamicallyLegalDialect<func::FuncDialect>(
         [](Operation *op) { return isLegal(op); });
     target.addIllegalDialect<memref::MemRefDialect>();
+    target.addLegalDialect<emitc::EmitCDialect>();
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
@@ -106,6 +135,8 @@ void mlir::populateMemRefToEmitCConversionPatterns(RewritePatternSet &patterns,
                                                                  converter);
   populateCallOpTypeConversionPattern(patterns, converter);
   populateReturnOpTypeConversionPattern(patterns, converter);
+  patterns.add<ConvertLoad>(converter, patterns.getContext());
+  patterns.add<ConvertStore>(converter, patterns.getContext());
 }
 
 std::unique_ptr<OperationPass<>> mlir::createConvertMemRefToEmitCPass() {

diff --git a/mlir/test/Conversion/MemRefToEmitC/memref-to-emitc.mlir b/mlir/test/Conversion/MemRefToEmitC/memref-to-emitc.mlir
@@ -22,3 +22,17 @@ func.func @memref_call(%arg0 : memref<32xf32>) {
   func.call @memref_return(%arg0) : (memref<32xf32>) -> memref<32xf32>
   func.return
 }
+
+// -----
+
+// CHECK-LABEL: memref_load_store
+// CHECK-SAME:  %[[arg0:.*]]: !emitc.array<4x8xf32>, %[[arg1:.*]]: !emitc.array<3x5xf32>
+// CHECK-SAME:  %[[i:.*]]: index, %[[j:.*]]: index
+func.func @memref_load_store(%in: memref<4x8xf32>, %out: memref<3x5xf32>, %i: index, %j: index) {
+  // CHECK: %[[load:.*]] = emitc.subscript %[[arg0]][%[[i]], %[[j]]] : <4x8xf32>
+  %0 = memref.load %in[%i, %j] : memref<4x8xf32>
+  // CHECK: %[[store_loc:.*]] = emitc.subscript %[[arg1]][%[[i]], %[[j]]] : <3x5xf32>
+  // CHECK: emitc.assign %[[load]] : f32 to %[[store_loc:.*]] : f32
+  memref.store %0, %out[%i, %j] : memref<3x5xf32>
+  return
+}