[Mlir-commits] [mlir] [mlir][EmitC]Expand the MemRefToEmitC pass - Lowering `AllocOp` (PR #148257)

[Jaden Angella]

================
@@ -77,6 +81,95 @@ struct ConvertAlloca final : public OpConversionPattern<memref::AllocaOp> {
   }
 };
 
+struct ConvertAlloc final : public OpConversionPattern<memref::AllocOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(memref::AllocOp allocOp, OpAdaptor operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    mlir::Location loc = allocOp.getLoc();
+    MemRefType memrefType = allocOp.getType();
+    if (!memrefType.hasStaticShape()) {
+      // TODO: Handle Dynamic shapes in the future. If the size
+      // of the allocation is the result of some function, we could
+      // potentially evaluate the function and use the result in the call to
+      // allocate.
+      return rewriter.notifyMatchFailure(
+          loc, "cannot transform alloc with dynamic shape");
+    }
+
+    mlir::Type sizeTType = mlir::emitc::SizeTType::get(rewriter.getContext());
+    Type elementType = memrefType.getElementType();
+    mlir::emitc::CallOpaqueOp sizeofElementOp =
+        rewriter.create<mlir::emitc::CallOpaqueOp>(
+            loc, sizeTType, rewriter.getStringAttr("sizeof"),
+            mlir::ValueRange{},
+            mlir::ArrayAttr::get(rewriter.getContext(),
+                                 {mlir::TypeAttr::get(elementType)}));
+    mlir::Value sizeofElement = sizeofElementOp.getResult(0);
+
+    unsigned int elementWidth = elementType.getIntOrFloatBitWidth();
+    IntegerAttr indexAttr = rewriter.getIndexAttr(elementWidth);
+
+    mlir::Value numElements;
+    numElements = rewriter.create<emitc::ConstantOp>(
+        loc, rewriter.getIndexType(), indexAttr);
+    mlir::Value totalSizeBytes = rewriter.create<emitc::MulOp>(
+        loc, sizeTType, sizeofElement, numElements);
+
+    emitc::CallOpaqueOp allocCall;
+    StringAttr allocFunctionName = allocOp.getAlignment()
+                                       ? rewriter.getStringAttr("aligned_alloc")
+                                       : rewriter.getStringAttr("malloc");
+    mlir::Value alignmentValue;
+    if (allocOp.getAlignment()) {
+      alignmentValue = rewriter.create<emitc::ConstantOp>(
+          loc, sizeTType,
+          rewriter.getIntegerAttr(
+              rewriter.getIndexType(),
+              alignedAllocationGetAlignment(allocOp, elementWidth)));
+    }
+
+    SmallVector<mlir::Value, 2> argsVec;
+    if (allocOp.getAlignment())
+      argsVec.push_back(alignmentValue);
+    argsVec.push_back(totalSizeBytes);
+    mlir::ValueRange args(argsVec);
+
+    allocCall = rewriter.create<emitc::CallOpaqueOp>(
+        loc,
+        emitc::PointerType::get(
+            mlir::emitc::OpaqueType::get(rewriter.getContext(), "void")),
+        allocFunctionName, args);
+
+    emitc::PointerType targetPointerType = emitc::PointerType::get(elementType);
+    emitc::CastOp castOp = rewriter.create<emitc::CastOp>(
+        loc, targetPointerType, allocCall.getResult(0));
+
+    rewriter.replaceOp(allocOp, castOp);
+    return success();
+  }
+
+  /// The minimum alignment to use with aligned_alloc (has to be a power of 2).
+  static constexpr uint64_t kMinAlignedAllocAlignment = 16UL;
+
+  /// Computes the alignment for aligned_alloc used to allocate the buffer for
+  /// the memory allocation op.
+  ///
+  /// Aligned_alloc requires the allocation size to be a power of two, and the
+  /// allocation size to be a multiple of the alignment.
+  int64_t alignedAllocationGetAlignment(memref::AllocOp op,
+                                        unsigned int elementWidth) const {
+    if (std::optional<uint64_t> alignment = op.getAlignment())
----------------
Jaddyen wrote:

This is true!
I can simply get the alignment without further checks

https://github.com/llvm/llvm-project/pull/148257