[MLIR][StandardToLLVM] Move `copyUnrankedDescriptors` to pattern

[frgossen]

Summary

Make the function copyUnrankedDescriptors accessible in
ConvertToLLVMPattern.

[ftynse]

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

122

How does it know where the memory was?

[frgossen]

Address comment

[frgossen]

Thanks!

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

122

It decides that based on toDynamic. Clarified this in the comment.

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

Value getNumElements(Location loc, ArrayRef<Value> shape,
                        ConversionPatternRewriter &rewriter) const;
 
   /// Creates and populates a canonical memref descriptor struct.
   MemRefDescriptor
   createMemRefDescriptor(Location loc, MemRefType memRefType,
                          Value allocatedPtr, Value alignedPtr,
                          ArrayRef<Value> sizes, ArrayRef<Value> strides,
                          ConversionPatternRewriter &rewriter) const;
+
+  /// Copies the memory descriptor for any operands that were unranked
+  /// descriptors originally to heap-allocated memory (if toDynamic is true) or
+  /// to stack-allocated memory (otherwise). Also frees the previously used
+  /// memory if it was not on the stack.
+  LogicalResult copyUnrankedDescriptors(OpBuilder &builder, Location loc,
+                                        TypeRange origTypes,
+                                        SmallVectorImpl<Value> &operands,
+                                        bool toDynamic) const;
 };
 
 /// Utility class for operation conversions targeting the LLVM dialect that
 /// match exactly one source operation.
 template <typename SourceOp>
 class ConvertOpToLLVMPattern : public ConvertToLLVMPattern {
 public:
   explicit ConvertOpToLLVMPattern(LLVMTypeConverter &typeConverter,

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

 //===- Pattern.cpp - Conversion pattern to the LLVM dialect ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/AffineMap.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // ConvertToLLVMPattern
MemRefDescriptor ConvertToLLVMPattern::createMemRefDescriptor(
 
   // Field 5: Strides.
   for (auto en : llvm::enumerate(strides))
     memRefDescriptor.setStride(rewriter, loc, en.index(), en.value());
 
   return memRefDescriptor;
 }
 
+LogicalResult ConvertToLLVMPattern::copyUnrankedDescriptors(
+    OpBuilder &builder, Location loc, TypeRange origTypes,
+    SmallVectorImpl<Value> &operands, bool toDynamic) const {
+  assert(origTypes.size() == operands.size() &&
+         "expected as may original types as operands");
+
+  // Find operands of unranked memref type and store them.
+  SmallVector<UnrankedMemRefDescriptor, 4> unrankedMemrefs;
+  for (unsigned i = 0, e = operands.size(); i < e; ++i)
+    if (origTypes[i].isa<UnrankedMemRefType>())
+      unrankedMemrefs.emplace_back(operands[i]);
+
+  if (unrankedMemrefs.empty())
+    return success();
+
+  // Compute allocation sizes.
+  SmallVector<Value, 4> sizes;
+  UnrankedMemRefDescriptor::computeSizes(builder, loc, *getTypeConverter(),
+                                         unrankedMemrefs, sizes);
+
+  // Get frequently used types.
+  MLIRContext *context = builder.getContext();
+  Type voidPtrType = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
+  auto i1Type = IntegerType::get(context, 1);
+  Type indexType = getTypeConverter()->getIndexType();
+
+  // Find the malloc and free, or declare them if necessary.
+  auto module = builder.getInsertionPoint()->getParentOfType<ModuleOp>();
+  LLVM::LLVMFuncOp freeFunc, mallocFunc;
+  if (toDynamic)
+    mallocFunc = LLVM::lookupOrCreateMallocFn(module, indexType);
+  if (!toDynamic)
+    freeFunc = LLVM::lookupOrCreateFreeFn(module);
+
+  // Initialize shared constants.
+  Value zero =
+      builder.create<LLVM::ConstantOp>(loc, i1Type, builder.getBoolAttr(false));
+
+  unsigned unrankedMemrefPos = 0;
+  for (unsigned i = 0, e = operands.size(); i < e; ++i) {
+    Type type = origTypes[i];
+    if (!type.isa<UnrankedMemRefType>())
+      continue;
+    Value allocationSize = sizes[unrankedMemrefPos++];
+    UnrankedMemRefDescriptor desc(operands[i]);
+
+    // Allocate memory, copy, and free the source if necessary.
+    Value memory =
+        toDynamic
+            ? builder.create<LLVM::CallOp>(loc, mallocFunc, allocationSize)
+                  .getResult(0)
+            : builder.create<LLVM::AllocaOp>(loc, voidPtrType, allocationSize,
+                                             /*alignment=*/0);
+    Value source = desc.memRefDescPtr(builder, loc);
+    builder.create<LLVM::MemcpyOp>(loc, memory, source, allocationSize, zero);
+    if (!toDynamic)
+      builder.create<LLVM::CallOp>(loc, freeFunc, source);
+
+    // Create a new descriptor. The same descriptor can be returned multiple
+    // times, attempting to modify its pointer can lead to memory leaks
+    // (allocated twice and overwritten) or double frees (the caller does not
+    // know if the descriptor points to the same memory).
+    Type descriptorType = getTypeConverter()->convertType(type);
+    if (!descriptorType)
+      return failure();
+    auto updatedDesc =
+        UnrankedMemRefDescriptor::undef(builder, loc, descriptorType);
+    Value rank = desc.rank(builder, loc);
+    updatedDesc.setRank(builder, loc, rank);
+    updatedDesc.setMemRefDescPtr(builder, loc, memory);
+
+    operands[i] = updatedDesc;
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Detail methods
 //===----------------------------------------------------------------------===//
 
 /// Replaces the given operation "op" with a new operation of type "targetOp"
 /// and given operands.
 LogicalResult LLVM::detail::oneToOneRewrite(
     Operation *op, StringRef targetOp, ValueRange operands,

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

matchAndRewrite(memref::AllocaScopeOp allocaScopeOp, ArrayRef<Value> operands,
 
     // Replace the op with values return from the body region.
     rewriter.replaceOp(allocaScopeOp, continueBlock->getArguments());
 
     return success();
   }
 };
 
-/// Copies the shaped descriptor part to (if `toDynamic` is set) or from
-/// (otherwise) the dynamically allocated memory for any operands that were
-/// unranked descriptors originally.
-static LogicalResult copyUnrankedDescriptors(OpBuilder &builder, Location loc,
-                                             LLVMTypeConverter &typeConverter,
-                                             TypeRange origTypes,
-                                             SmallVectorImpl<Value> &operands,
-                                             bool toDynamic) {
-  assert(origTypes.size() == operands.size() &&
-         "expected as may original types as operands");
-
-  // Find operands of unranked memref type and store them.
-  SmallVector<UnrankedMemRefDescriptor, 4> unrankedMemrefs;
-  for (unsigned i = 0, e = operands.size(); i < e; ++i)
-    if (origTypes[i].isa<UnrankedMemRefType>())
-      unrankedMemrefs.emplace_back(operands[i]);
-
-  if (unrankedMemrefs.empty())
-    return success();
-
-  // Compute allocation sizes.
-  SmallVector<Value, 4> sizes;
-  UnrankedMemRefDescriptor::computeSizes(builder, loc, typeConverter,
-                                         unrankedMemrefs, sizes);
-
-  // Get frequently used types.
-  MLIRContext *context = builder.getContext();
-  Type voidPtrType = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
-  auto i1Type = IntegerType::get(context, 1);
-  Type indexType = typeConverter.getIndexType();
-
-  // Find the malloc and free, or declare them if necessary.
-  auto module = builder.getInsertionPoint()->getParentOfType<ModuleOp>();
-  LLVM::LLVMFuncOp freeFunc, mallocFunc;
-  if (toDynamic)
-    mallocFunc = LLVM::lookupOrCreateMallocFn(module, indexType);
-  if (!toDynamic)
-    freeFunc = LLVM::lookupOrCreateFreeFn(module);
-
-  // Initialize shared constants.
-  Value zero =
-      builder.create<LLVM::ConstantOp>(loc, i1Type, builder.getBoolAttr(false));
-
-  unsigned unrankedMemrefPos = 0;
-  for (unsigned i = 0, e = operands.size(); i < e; ++i) {
-    Type type = origTypes[i];
-    if (!type.isa<UnrankedMemRefType>())
-      continue;
-    Value allocationSize = sizes[unrankedMemrefPos++];
-    UnrankedMemRefDescriptor desc(operands[i]);
-
-    // Allocate memory, copy, and free the source if necessary.
-    Value memory =
-        toDynamic
-            ? builder.create<LLVM::CallOp>(loc, mallocFunc, allocationSize)
-                  .getResult(0)
-            : builder.create<LLVM::AllocaOp>(loc, voidPtrType, allocationSize,
-                                             /*alignment=*/0);
-
-    Value source = desc.memRefDescPtr(builder, loc);
-    builder.create<LLVM::MemcpyOp>(loc, memory, source, allocationSize, zero);
-    if (!toDynamic)
-      builder.create<LLVM::CallOp>(loc, freeFunc, source);
-
-    // Create a new descriptor. The same descriptor can be returned multiple
-    // times, attempting to modify its pointer can lead to memory leaks
-    // (allocated twice and overwritten) or double frees (the caller does not
-    // know if the descriptor points to the same memory).
-    Type descriptorType = typeConverter.convertType(type);
-    if (!descriptorType)
-      return failure();
-    auto updatedDesc =
-        UnrankedMemRefDescriptor::undef(builder, loc, descriptorType);
-    Value rank = desc.rank(builder, loc);
-    updatedDesc.setRank(builder, loc, rank);
-    updatedDesc.setMemRefDescPtr(builder, loc, memory);
-
-    operands[i] = updatedDesc;
-  }
-
-  return success();
-}
-
 // A CallOp automatically promotes MemRefType to a sequence of alloca/store and
 // passes the pointer to the MemRef across function boundaries.
 template <typename CallOpType>
 struct CallOpInterfaceLowering : public ConvertOpToLLVMPattern<CallOpType> {
   using ConvertOpToLLVMPattern<CallOpType>::ConvertOpToLLVMPattern;
   using Super = CallOpInterfaceLowering<CallOpType>;
   using Base = ConvertOpToLLVMPattern<CallOpType>;
 
matchAndRewrite(CallOpType callOp, ArrayRef<Value> operands,
 
     if (this->getTypeConverter()->getOptions().useBarePtrCallConv) {
       // For the bare-ptr calling convention, promote memref results to
       // descriptors.
       assert(results.size() == resultTypes.size() &&
              "The number of arguments and types doesn't match");
       this->getTypeConverter()->promoteBarePtrsToDescriptors(
           rewriter, callOp.getLoc(), resultTypes, results);
-    } else if (failed(copyUnrankedDescriptors(rewriter, callOp.getLoc(),
-                                              *this->getTypeConverter(),
+    } else if (failed(this->copyUnrankedDescriptors(rewriter, callOp.getLoc(),
                                                     resultTypes, results,
                                                     /*toDynamic=*/false))) {
       return failure();
     }
 
     rewriter.replaceOp(callOp, results);
     return success();
   }
 };
 
if (getTypeConverter()->getOptions().useBarePtrCallConv) {
           // Unranked memref is not supported in the bare pointer calling
           // convention.
           return failure();
         }
         updatedOperands.push_back(newOperand);
       }
     } else {
       updatedOperands = llvm::to_vector<4>(operands);
-      (void)copyUnrankedDescriptors(rewriter, loc, *getTypeConverter(),
-                                    op.getOperands().getTypes(),
+      (void)copyUnrankedDescriptors(rewriter, loc, op.getOperands().getTypes(),
                                     updatedOperands,
                                     /*toDynamic=*/true);
     }
 
     // If ReturnOp has 0 or 1 operand, create it and return immediately.
     if (numArguments == 0) {
       rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, TypeRange(), ValueRange(),
                                                   op->getAttrs());