[MLIR] Check for static shape before bare pointer conversion

[chelini]

Summary

Originally in the returnOp conversion, the result type was changing to bare
pointer if the type was a memref. This is incorrect as conversion to bare
pointer can only be done if the memref has static shape, strides and offset.

[ftynse]

Thanks!

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

public:
 
   /// Returns the size of the memref descriptor object in bytes.
   unsigned getMemRefDescriptorSize(MemRefType type, const DataLayout &layout);
 
   /// Returns the size of the unranked memref descriptor object in bytes.
   unsigned getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
                                            const DataLayout &layout);
 
+  /// Check if a memref type can be converted to a bare pointer.
+  bool canConvertToBarePtr(BaseMemRefType type);
+
 protected:
   /// Pointer to the LLVM dialect.
   LLVM::LLVMDialect *llvmDialect;
 
 private:
   /// Convert a function type.  The arguments and results are converted one by
   /// one.  Additionally, if the function returns more than one value, pack the
   /// results into an LLVM IR structure type so that the converted function type
private:
   /// - an integer rank, followed by
   /// - a pointer to the memref descriptor struct.
   /// For example, memref<*xf32> is converted to the following list:
   /// i64 (rank)
   /// !llvm<"i8*"> (type-erased pointer).
   /// These types can be recomposed to a unranked memref descriptor struct.
   SmallVector<Type, 2> getUnrankedMemRefDescriptorFields();
 
-  // Convert an unranked memref type to an LLVM type that captures the
-  // runtime rank and a pointer to the static ranked memref desc
+  /// Convert an unranked memref type to an LLVM type that captures the
+  /// runtime rank and a pointer to the static ranked memref desc
   Type convertUnrankedMemRefType(UnrankedMemRefType type);
 
   /// Convert a memref type to a bare pointer to the memref element type.
   Type convertMemRefToBarePtr(BaseMemRefType type);
 
   /// Convert a 1D vector type into an LLVM vector type.
   Type convertVectorType(VectorType type);
 

mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp

matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
     SmallVector<Value, 4> updatedOperands;
 
     if (getTypeConverter()->getOptions().useBarePtrCallConv) {
       // For the bare-ptr calling convention, extract the aligned pointer to
       // be returned from the memref descriptor.
       for (auto it : llvm::zip(op->getOperands(), adaptor.getOperands())) {
         Type oldTy = std::get<0>(it).getType();
         Value newOperand = std::get<1>(it);
-        if (oldTy.isa<MemRefType>()) {
+        if (oldTy.isa<MemRefType>() && getTypeConverter()->canConvertToBarePtr(
+                                           oldTy.cast<BaseMemRefType>())) {
           MemRefDescriptor memrefDesc(newOperand);
           newOperand = memrefDesc.alignedPtr(rewriter, loc);
         } else if (oldTy.isa<UnrankedMemRefType>()) {
           // Unranked memref is not supported in the bare pointer calling
           // convention.
           return failure();
         }
         updatedOperands.push_back(newOperand);

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

 
 Type LLVMTypeConverter::convertUnrankedMemRefType(UnrankedMemRefType type) {
   if (!convertType(type.getElementType()))
     return {};
   return LLVM::LLVMStructType::getLiteral(&getContext(),
                                           getUnrankedMemRefDescriptorFields());
 }
 
-/// Convert a memref type to a bare pointer to the memref element type.
-Type LLVMTypeConverter::convertMemRefToBarePtr(BaseMemRefType type) {
+// Check if a memref type can be converted to a bare pointer.
+bool LLVMTypeConverter::canConvertToBarePtr(BaseMemRefType type) {
   if (type.isa<UnrankedMemRefType>())
     // Unranked memref is not supported in the bare pointer calling convention.
-    return {};
+    return false;
 
   // Check that the memref has static shape, strides and offset. Otherwise, it
   // cannot be lowered to a bare pointer.
   auto memrefTy = type.cast<MemRefType>();
   if (!memrefTy.hasStaticShape())
-    return {};
+    return false;
 
   int64_t offset = 0;
   SmallVector<int64_t, 4> strides;
   if (failed(getStridesAndOffset(memrefTy, strides, offset)))
-    return {};
+    return false;
 
   for (int64_t stride : strides)
     if (ShapedType::isDynamicStrideOrOffset(stride))
-      return {};
+      return false;
 
   if (ShapedType::isDynamicStrideOrOffset(offset))
-    return {};
+    return false;
+
+  return true;
+}
 
+/// Convert a memref type to a bare pointer to the memref element type.
+Type LLVMTypeConverter::convertMemRefToBarePtr(BaseMemRefType type) {
+  if (!canConvertToBarePtr(type))
+    return {};
   Type elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
   return LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
 }
 
 /// Convert an n-D vector type to an LLVM vector type:
 ///  * 0-D `vector<T>` are converted to vector<1xT>

mlir/test/Conversion/FuncToLLVM/func-memref-return.mlir

func @check_memref_func_call(%in : memref<10xi8>) -> memref<20xi8> {
   // BAREPTR-NEXT:    %[[desc6:.*]] = llvm.insertvalue %[[c20]], %[[desc4]][3, 0] : !llvm.struct<(ptr<i8>, ptr<i8>, i64, array<1 x i64>, array<1 x i64>)>
   // BAREPTR-NEXT:    %[[c1:.*]] = llvm.mlir.constant(1 : index) : i64
   // BAREPTR-NEXT:    %[[outDesc:.*]] = llvm.insertvalue %[[c1]], %[[desc6]][4, 0] : !llvm.struct<(ptr<i8>, ptr<i8>, i64, array<1 x i64>, array<1 x i64>)>
   %res = call @foo(%in) : (memref<10xi8>) -> (memref<20xi8>)
   // BAREPTR-NEXT:    %[[res:.*]] = llvm.extractvalue %[[outDesc]][1] : !llvm.struct<(ptr<i8>, ptr<i8>, i64, array<1 x i64>, array<1 x i64>)>
   // BAREPTR-NEXT:    llvm.return %[[res]] : !llvm.ptr<i8>
   return %res : memref<20xi8>
 }
+
+// -----
+
+// BAREPTR-LABEL: func @check_return(
+// BAREPTR-SAME: %{{.*}}: memref<?xi8>) -> memref<?xi8> 
+func @check_return(%in : memref<?xi8>) -> memref<?xi8> {
+  // BAREPTR: llvm.return {{.*}} : !llvm.struct<(ptr<i8>, ptr<i8>, i64, array<1 x i64>, array<1 x i64>)>
+  return %in : memref<?xi8>
+}