[mlir] NFC: Expose `getElementPtrType` and `getSizes` methods of AllocOpLowering.

[pifon2a]

Summary

These methods are very useful outside of AllocLikeOpLowering.

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

935

I would prefer for this method to be split into two pieces

getMemRefDescriptorSizes that gives you the SmallVector of sizes.

and

getMemRefDescriptorSizeInBytes that takes the sizes and gives you the Value for the size.

Longer term, we could then reuse the former in more places.

938

Making this one reused makes sense for a private function but not for a public API I think.

[pifon2a]

Split getSizes() into getMemRefDescriptorsizes and getCumulativeSizeInBytes.

[herhut]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

943

Is this actually correct? If the memref is 0d, it should have no sizes. I think this was convenient for the later computation of the byte size only.

947

Should this just return the value and create its own one? I think it is better to not share the one constant. They are cheap and cse will fix it. And even if not, they don't turn into code anyway.

[pifon2a]

Address the comments.

[herhut]

Please fix linter error and wait for @ftynse to comment.

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1960

Accidental edit?

[pifon2a]

Fix linter.

[ftynse]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

941

Nit: can this use ShapedType::kDynamicSize instead of hardcoding -1 ?

961

Nit: template seems no longer necessary

[pifon2a]

Address the comments.

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

Value getStridedElementPtr(Location loc, Type elementTypePtr,
                              Value descriptor, ValueRange indices,
                              ArrayRef<int64_t> strides, int64_t offset,
                              ConversionPatternRewriter &rewriter) const;
 
   Value getDataPtr(Location loc, MemRefType type, Value memRefDesc,
                    ValueRange indices, ConversionPatternRewriter &rewriter,
                    llvm::Module &module) const;
 
+  /// Returns the type of a pointer to an element of the memref.
+  Type getElementPtrType(MemRefType type) const;
+
+  /// Determines sizes to be used in the memref descriptor.
+  void getMemRefDescriptorSizes(Location loc, MemRefType memRefType,
+                                ArrayRef<Value> dyn_sizes,
+                                ConversionPatternRewriter &rewriter,
+                                SmallVectorImpl<Value> &sizes) const;
+
+  /// Computes total size in bytes of to store the given shape.
+  void getCumulativeSizeInBytes(Location loc, Type elementType,
+                                ArrayRef<Value> shape,
+                                ConversionPatternRewriter &rewriter,
+                                Value &cumulativeSizeInBytes, Value &one) const;
+
 protected:
   /// Reference to the type converter, with potential extensions.
   LLVMTypeConverter &typeConverter;
 
   /// Reference to the llvm lowering options.
   const LowerToLLVMOptions &options;
 };
 

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

Value ConvertToLLVMPattern::getDataPtr(Location loc, MemRefType type,
   SmallVector<int64_t, 4> strides;
   auto successStrides = getStridesAndOffset(type, strides, offset);
   assert(succeeded(successStrides) && "unexpected non-strided memref");
   (void)successStrides;
   return getStridedElementPtr(loc, ptrType, memRefDesc, indices, strides,
                               offset, rewriter);
 }
 
+Type ConvertToLLVMPattern::getElementPtrType(MemRefType type) const {
+  auto elementType = type.getElementType();
+  auto structElementType = typeConverter.convertType(elementType);
+  return structElementType.cast<LLVM::LLVMType>().getPointerTo(
+      type.getMemorySpace());
+}
+
+void ConvertToLLVMPattern::getMemRefDescriptorSizes(
+    Location loc, MemRefType memRefType, ArrayRef<Value> dyn_sizes,
+    ConversionPatternRewriter &rewriter, SmallVectorImpl<Value> &sizes) const {
+  sizes.reserve(memRefType.getRank());
+  unsigned i = 0;
+  for (int64_t s : memRefType.getShape())
+    sizes.push_back(s == -1 ? dyn_sizes[i++]
+                            : createIndexConstant(rewriter, loc, s));
+  if (sizes.empty())
+    sizes.push_back(createIndexConstant(rewriter, loc, 1));
+}
+
+void ConvertToLLVMPattern::getCumulativeSizeInBytes(
+    Location loc, Type elementType, ArrayRef<Value> sizes,
+    ConversionPatternRewriter &rewriter, Value &cumulativeSizeInBytes,
+    Value &one) const {
+  // Compute the total number of memref elements.
+  cumulativeSizeInBytes = sizes.front();
+  for (unsigned i = 1, e = sizes.size(); i < e; ++i)
+    cumulativeSizeInBytes = rewriter.create<LLVM::MulOp>(
+        loc, getIndexType(), ArrayRef<Value>{cumulativeSizeInBytes, sizes[i]});
+
+  // Compute the size of an individual element. This emits the MLIR equivalent
+  // of the following sizeof(...) implementation in LLVM IR:
+  //   %0 = getelementptr %elementType* null, %indexType 1
+  //   %1 = ptrtoint %elementType* %0 to %indexType
+  // which is a common pattern of getting the size of a type in bytes.
+  auto convertedPtrType = typeConverter.convertType(elementType)
+                              .template cast<LLVM::LLVMType>()
+                              .getPointerTo();
+  auto nullPtr = rewriter.create<LLVM::NullOp>(loc, convertedPtrType);
+  one = createIndexConstant(rewriter, loc, 1);
+  auto gep = rewriter.create<LLVM::GEPOp>(loc, convertedPtrType,
+                                          ArrayRef<Value>{nullPtr, one});
+  auto elementSize =
+      rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gep);
+  cumulativeSizeInBytes = rewriter.create<LLVM::MulOp>(
+      loc, getIndexType(), ArrayRef<Value>{cumulativeSizeInBytes, elementSize});
+}
+
 /// Only retain those attributes that are not constructed by
 /// `LLVMFuncOp::build`. If `filterArgAttrs` is set, also filter out argument
 /// attributes.
 static void filterFuncAttributes(ArrayRef<NamedAttribute> attrs,
                                  bool filterArgAttrs,
                                  SmallVectorImpl<NamedAttribute> &result) {
   for (const auto &attr : attrs) {
     if (attr.first == SymbolTable::getSymbolAttrName() ||
struct AllocLikeOpLowering : public ConvertOpToLLVMPattern<AllocLikeOp> {
   /// Creates and populates the memref descriptor struct given all its fields.
   /// This method also performs any post allocation alignment needed for heap
   /// allocations when `accessAlignment` is non null. This is used with
   /// allocators that do not support alignment.
   MemRefDescriptor createMemRefDescriptor(
       Location loc, ConversionPatternRewriter &rewriter, MemRefType memRefType,
       Value allocatedTypePtr, Value allocatedBytePtr, Value accessAlignment,
       uint64_t offset, ArrayRef<int64_t> strides, ArrayRef<Value> sizes) const {
-    auto elementPtrType = getElementPtrType(memRefType);
+    auto elementPtrType = this->getElementPtrType(memRefType);
     auto structType = typeConverter.convertType(memRefType);
     auto memRefDescriptor = MemRefDescriptor::undef(rewriter, loc, structType);
 
     // Field 1: Allocated pointer, used for malloc/free.
     memRefDescriptor.setAllocatedPtr(rewriter, loc, allocatedTypePtr);
 
     // Field 2: Actual aligned pointer to payload.
     Value alignedBytePtr = allocatedTypePtr;
MemRefDescriptor createMemRefDescriptor(
     for (auto indexedSize : llvm::enumerate(sizes)) {
       int64_t index = indexedSize.index();
       memRefDescriptor.setSize(rewriter, loc, index, indexedSize.value());
       memRefDescriptor.setStride(rewriter, loc, index, strideValues[index]);
     }
     return memRefDescriptor;
   }
 
-  /// Determines sizes to be used in the memref descriptor.
-  void getSizes(Location loc, MemRefType memRefType, ArrayRef<Value> operands,
-                ConversionPatternRewriter &rewriter,
-                SmallVectorImpl<Value> &sizes, Value &cumulativeSize,
-                Value &one) const {
-    sizes.reserve(memRefType.getRank());
-    unsigned i = 0;
-    for (int64_t s : memRefType.getShape())
-      sizes.push_back(s == -1 ? operands[i++]
-                              : createIndexConstant(rewriter, loc, s));
-    if (sizes.empty())
-      sizes.push_back(createIndexConstant(rewriter, loc, 1));
-
-    // Compute the total number of memref elements.
-    cumulativeSize = sizes.front();
-    for (unsigned i = 1, e = sizes.size(); i < e; ++i)
-      cumulativeSize = rewriter.create<LLVM::MulOp>(
-          loc, getIndexType(), ArrayRef<Value>{cumulativeSize, sizes[i]});
-
-    // Compute the size of an individual element. This emits the MLIR equivalent
-    // of the following sizeof(...) implementation in LLVM IR:
-    //   %0 = getelementptr %elementType* null, %indexType 1
-    //   %1 = ptrtoint %elementType* %0 to %indexType
-    // which is a common pattern of getting the size of a type in bytes.
-    auto elementType = memRefType.getElementType();
-    auto convertedPtrType = typeConverter.convertType(elementType)
-                                .template cast<LLVM::LLVMType>()
-                                .getPointerTo();
-    auto nullPtr = rewriter.create<LLVM::NullOp>(loc, convertedPtrType);
-    one = createIndexConstant(rewriter, loc, 1);
-    auto gep = rewriter.create<LLVM::GEPOp>(loc, convertedPtrType,
-                                            ArrayRef<Value>{nullPtr, one});
-    auto elementSize =
-        rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gep);
-    cumulativeSize = rewriter.create<LLVM::MulOp>(
-        loc, getIndexType(), ArrayRef<Value>{cumulativeSize, elementSize});
-  }
-
-  /// Returns the type of a pointer to an element of the memref.
-  Type getElementPtrType(MemRefType memRefType) const {
-    auto elementType = memRefType.getElementType();
-    auto structElementType = typeConverter.convertType(elementType);
-    return structElementType.template cast<LLVM::LLVMType>().getPointerTo(
-        memRefType.getMemorySpace());
-  }
-
   /// Returns the memref's element size in bytes.
   // TODO: there are other places where this is used. Expose publicly?
   static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
     auto elementType = memRefType.getElementType();
 
     unsigned sizeInBits;
     if (elementType.isIntOrFloat()) {
       sizeInBits = elementType.getIntOrFloatBitWidth();
struct AllocLikeOpLowering : public ConvertOpToLLVMPattern<AllocLikeOp> {
 
   /// Allocates the underlying buffer using the right call. `allocatedBytePtr`
   /// is set to null for stack allocations. `accessAlignment` is set if
   /// alignment is needed post allocation (for eg. in conjunction with malloc).
   Value allocateBuffer(Location loc, Value cumulativeSize, Operation *op,
                        MemRefType memRefType, Value one, Value &accessAlignment,
                        Value &allocatedBytePtr,
                        ConversionPatternRewriter &rewriter) const {
-    auto elementPtrType = getElementPtrType(memRefType);
+    auto elementPtrType = this->getElementPtrType(memRefType);
 
     // With alloca, one gets a pointer to the element type right away.
     // For stack allocations.
     if (auto allocaOp = dyn_cast<AllocaOp>(op)) {
       allocatedBytePtr = nullptr;
       accessAlignment = nullptr;
       return rewriter.create<LLVM::AllocaOp>(
           loc, elementPtrType, cumulativeSize,
struct AllocLikeOpLowering : public ConvertOpToLLVMPattern<AllocLikeOp> {
   // An `alloca` is converted into a definition of a memref descriptor value and
   // an llvm.alloca to allocate the underlying data buffer.
   void rewrite(Operation *op, ArrayRef<Value> operands,
                ConversionPatternRewriter &rewriter) const override {
     MemRefType memRefType = cast<AllocLikeOp>(op).getType();
     auto loc = op->getLoc();
 
     // Get actual sizes of the memref as values: static sizes are constant
+    //
     // values and dynamic sizes are passed to 'alloc' as operands.  In case of
     // zero-dimensional memref, assume a scalar (size 1).
     SmallVector<Value, 4> sizes;
-    Value cumulativeSize, one;
-    getSizes(loc, memRefType, operands, rewriter, sizes, cumulativeSize, one);
+    this->getMemRefDescriptorSizes(loc, memRefType, operands, rewriter, sizes);
 
+    Value cumulativeSize, one;
+    this->getCumulativeSizeInBytes(loc, memRefType.getElementType(), sizes,
+                                   rewriter, cumulativeSize, one);
     // Allocate the underlying buffer.
     // Value holding the alignment that has to be performed post allocation
     // (in conjunction with allocators that do not support alignment, eg.
     // malloc); nullptr if no such adjustment needs to be performed.
     Value accessAlignment;
     // Byte pointer to the allocated buffer.
     Value allocatedBytePtr;
     Value allocatedTypePtr =