[mlir] Require std.alloc() ops to have canonical layout during LLVM lowering.

[csigg]

Summary

The current code allows strided layouts, but the number of elements allocated is ambiguous. It could be either the number of elements in the shape (the current implementation), or the amount of elements required to not index out-of-bounds with the given maps (which would require evaluating the layout map).

If we require the canonical layouts, the two will be the same.

[mehdi_amini]

[mlir] Require std.alloc() ops to have canonical layout during LLVM lowering.

Do you mean identity layout? I thought the strided form to be canonical otherwise?

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 
 private:
   static MemRefType getMemRefResultType(Operation *op) {
     return op->getResult(0).getType().cast<MemRefType>();
   }
 
   LogicalResult match(Operation *op) const override {
     MemRefType memRefType = getMemRefResultType(op);
-    if (isSupportedMemRefType(memRefType))
-      return success();
-
-    int64_t offset;
-    SmallVector<int64_t, 4> strides;
-    if (failed(getStridesAndOffset(memRefType, strides, offset)))
-      return failure();
-
-    // Dynamic strides are ok if they can be deduced from dynamic sizes (which
-    // is guaranteed when getStridesAndOffset succeeded. Dynamic offset however
-    // can never be alloc'ed.
-    if (offset == MemRefType::getDynamicStrideOrOffset())
-      return failure();
-
-    return success();
+    return success(isSupportedMemRefType(memRefType));
   }
 
   // An `alloc` is converted into a definition of a memref descriptor value and
   // a call to `malloc` to allocate the underlying data buffer.  The memref
   // descriptor is of the LLVM structure type where:
   //   1. the first element is a pointer to the allocated (typed) data buffer,
   //   2. the second element is a pointer to the (typed) payload, aligned to the
   //      specified alignment,

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -convert-std-to-llvm -split-input-file -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: func @address_space(
 // CHECK-SAME:    !llvm.ptr<float, 7>
 func @address_space(%arg0 : memref<32xf32, affine_map<(d0) -> (d0)>, 7>) {
   %0 = alloc() : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   %1 = constant 7 : index
   // CHECK: llvm.load %{{.*}} : !llvm.ptr<float, 5>
   %2 = load %0[%1] : memref<32xf32, affine_map<(d0) -> (d0)>, 5>
   std.return
 }
 
-// CHECK-LABEL: func @strided_memref(
-func @strided_memref(%ind: index) {
-  %0 = alloc()[%ind] : memref<32x64xf32, affine_map<(i, j)[M] -> (32 + M * i + j)>>
-  std.return
-}
-
 // -----
 
 // CHECK-LABEL: func @rsqrt(
 // CHECK-SAME: !llvm.float
 func @rsqrt(%arg0 : f32) {
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1.000000e+00 : f32) : !llvm.float
   // CHECK: %[[SQRT:.*]] = "llvm.intr.sqrt"(%arg0) : (!llvm.float) -> !llvm.float
   // CHECK: %[[DIV:.*]] = llvm.fdiv %[[ONE]], %[[SQRT]] : !llvm.float