[mlir][MemRefToLLVM] Fix the lowering of extract_strided_metadata

[qcolombet]

Summary

The first result of the extract_strided_metadata operation is a MemRef, not a naked pointer.
This patch fixes the lowering of this operation in MemRefToLLVM so that we properly materialize the full MemRef structure and not just the base naked pointer.

[qcolombet]

• The order of evaluation of the arguments of a function is platform dependent, as a result the code that extracts the base and aligned pointers from the source memref would happen in different order on different platforms, hence the failure on debian. To fix that put the basePtr and alignedPtr in temporary variables before making the call.

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

public:
   static MemRefDescriptor undef(OpBuilder &builder, Location loc,
                                 Type descriptorType);
   /// Builds IR creating a MemRef descriptor that represents `type` and
   /// populates it with static shape and stride information extracted from the
   /// type.
   static MemRefDescriptor fromStaticShape(OpBuilder &builder, Location loc,
                                           LLVMTypeConverter &typeConverter,
                                           MemRefType type, Value memory);
+  static MemRefDescriptor fromStaticShape(OpBuilder &builder, Location loc,
+                                          LLVMTypeConverter &typeConverter,
+                                          MemRefType type, Value memory,
+                                          Value alignedMemory);
 
   /// Builds IR extracting the allocated pointer from the descriptor.
   Value allocatedPtr(OpBuilder &builder, Location loc);
   /// Builds IR inserting the allocated pointer into the descriptor.
   void setAllocatedPtr(OpBuilder &builder, Location loc, Value ptr);
 
   /// Builds IR extracting the aligned pointer from the descriptor.
   Value alignedPtr(OpBuilder &builder, Location loc);

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

 
 /// Builds IR creating a MemRef descriptor that represents `type` and
 /// populates it with static shape and stride information extracted from the
 /// type.
 MemRefDescriptor
 MemRefDescriptor::fromStaticShape(OpBuilder &builder, Location loc,
                                   LLVMTypeConverter &typeConverter,
                                   MemRefType type, Value memory) {
+  return fromStaticShape(builder, loc, typeConverter, type, memory, memory);
+}
+
+MemRefDescriptor MemRefDescriptor::fromStaticShape(
+    OpBuilder &builder, Location loc, LLVMTypeConverter &typeConverter,
+    MemRefType type, Value memory, Value alignedMemory) {
   assert(type.hasStaticShape() && "unexpected dynamic shape");
 
   // Extract all strides and offsets and verify they are static.
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   auto result = getStridesAndOffset(type, strides, offset);
   (void)result;
   assert(succeeded(result) && "unexpected failure in stride computation");
   assert(!ShapedType::isDynamicStrideOrOffset(offset) &&
          "expected static offset");
   assert(!llvm::any_of(strides, ShapedType::isDynamicStrideOrOffset) &&
          "expected static strides");
 
   auto convertedType = typeConverter.convertType(type);
   assert(convertedType && "unexpected failure in memref type conversion");
 
   auto descr = MemRefDescriptor::undef(builder, loc, convertedType);
   descr.setAllocatedPtr(builder, loc, memory);
-  descr.setAlignedPtr(builder, loc, memory);
+  descr.setAlignedPtr(builder, loc, alignedMemory);
   descr.setConstantOffset(builder, loc, offset);
 
   // Fill in sizes and strides
   for (unsigned i = 0, e = type.getRank(); i != e; ++i) {
     descr.setConstantSize(builder, loc, i, type.getDimSize(i));
     descr.setConstantStride(builder, loc, i, strides[i]);
   }
   return descr;

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

public:
   matchAndRewrite(memref::ExtractStridedMetadataOp extractStridedMetadataOp,
                   OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
 
     if (!LLVM::isCompatibleType(adaptor.getOperands().front().getType()))
       return failure();
 
     // Create the descriptor.
-    MemRefDescriptor sourceMemRef(adaptor.getOperands().front());
+    MemRefDescriptor sourceMemRef(adaptor.getSource());
     Location loc = extractStridedMetadataOp.getLoc();
     Value source = extractStridedMetadataOp.getSource();
 
     auto sourceMemRefType = source.getType().cast<MemRefType>();
     int64_t rank = sourceMemRefType.getRank();
     SmallVector<Value> results;
     results.reserve(2 + rank * 2);
 
     // Base buffer.
-    results.push_back(sourceMemRef.allocatedPtr(rewriter, loc));
+    Value baseBuffer = sourceMemRef.allocatedPtr(rewriter, loc);
+    Value alignedBuffer = sourceMemRef.alignedPtr(rewriter, loc);
+    MemRefDescriptor dstMemRef = MemRefDescriptor::fromStaticShape(
+        rewriter, loc, *getTypeConverter(),
+        extractStridedMetadataOp.getBaseBuffer().getType().cast<MemRefType>(),
+        baseBuffer, alignedBuffer);
+    results.push_back((Value)dstMemRef);
 
     // Offset.
     results.push_back(sourceMemRef.offset(rewriter, loc));
 
     // Sizes.
     for (unsigned i = 0; i < rank; ++i)
       results.push_back(sourceMemRef.size(rewriter, loc, i));
     // Strides.

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

 }
 
 // -----
 
 // CHECK-LABEL: func @extract_strided_metadata(
 // CHECK-SAME: %[[ARG:.*]]: memref
 // CHECK: %[[MEM_DESC:.*]] = builtin.unrealized_conversion_cast %[[ARG]] : memref<?x?xf32, strided<[?, ?], offset: ?>> to !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK: %[[BASE:.*]] = llvm.extractvalue %[[MEM_DESC]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK: %[[ALIGNED_BASE:.*]] = llvm.extractvalue %[[MEM_DESC]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK: %[[DESC:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+// CHECK: %[[DESC0:.*]] = llvm.insertvalue %[[BASE]], %[[DESC]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+// CHECK: %[[DESC1:.*]] = llvm.insertvalue %[[ALIGNED_BASE]], %[[DESC0]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
+// CHECK: %[[OFF0:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: %[[BASE_BUFFER_DESC:.*]] = llvm.insertvalue %[[OFF0]], %[[DESC1]][2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
 // CHECK: %[[OFFSET:.*]] = llvm.extractvalue %[[MEM_DESC]][2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK: %[[SIZE0:.*]] = llvm.extractvalue %[[MEM_DESC]][3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK: %[[SIZE1:.*]] = llvm.extractvalue %[[MEM_DESC]][3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK: %[[STRIDE0:.*]] = llvm.extractvalue %[[MEM_DESC]][4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 // CHECK: %[[STRIDE1:.*]] = llvm.extractvalue %[[MEM_DESC]][4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
 func.func @extract_strided_metadata(
     %ref: memref<?x?xf32, strided<[?,?], offset: ?>>) {
 
   %base, %offset, %sizes:2, %strides:2 =
     memref.extract_strided_metadata %ref : memref<?x?xf32, strided<[?,?], offset: ?>>
     -> memref<f32>, index,
        index, index,
        index, index
 
   return
 }