Enable ReassociatingReshapeOpConversion with "non-identity" layouts.

[Benoit]

Summary

Enable ReassociatingReshapeOpConversion with "non-identity" layouts.

This removes an early-return in this function, which seems unnecessary and is
preventing some memref.collapse_shape from converting to LLVM (see included lit test).

It seems unnecessary because the return message says "only empty layout map is supported"
but there actually is code in this function to deal with non-empty layout maps. Maybe
it refers to an earlier state of implementation and is just out of date?

Though, there is another concern about this early return: the condition that it actually
checks, {src,dst}MemrefType.getLayout().isIdentity(), is not quite the same as what the
return message says, "only empty layout map is supported". Stepping through this
getLayout().isIdentity() code in GDB, I found that it evaluates to .getAffineMap().isIdentity()
which does (AffineMap.cpp:271):

if (getNumDims() != getNumResults())
  return false;

This seems that it would always return false for memrefs of rank greater than 1 ?

[Benoit]

Rename the test function - was not match what it actually did

[nicolasvasilache]

This is due to the weird trifecta between:

1. memref without an affine_layout map being the convention to represent the canonical contiguous buffer of size equal to "the product of sizes"
2. memref with an identity affine_map layout being by convention equivalent to no affine_layout
3. memref with a strided layout that are currently represented by an affine_map.

This also intersects with the fact that a memref layout is provisioned to allow an arbitrarily complex composition of affine maps that I won't even get into.

I forget between the equivalence between 1. and 2. predates the strided memref or not (I think so but I am not 100% sure).

Things will get simpler once strided memref get back to using a different layout attribute offset:x, strides:[x,x,x] as originally intended.
However there are still many skeletons left because in the absence of a richer type/attribute system, it is unknown how the x in the strides relate to the ? in the sizes.
Properly encoding that relationship is key to being airtight in representing contiguous / non-contiguous layouts and reshapes with dynamic shapes.

Based on this info, please feel free to update docs so that it is less surprising to your future self and to others who will approach from the same angle as you did.

Thanks much!

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

1171

See my longer comment but please only relax this in the fully static case.

in the absence of a richer type/attribute system, it is unknown how the `x` in the strides relate to the `?` in the sizes.

[Benoit]

restrict to static shape case

[Benoit]

Thanks for the review and explanation. I have updated the code to apply your inline comment i.e. restrict the change to the static shape case.
I haven't updated docs as you suggested because I'm out of my depth here, sorry.

[mravishankar]

In D114808#3169700, @nicolasvasilache wrote:

This is due to the weird trifecta between:

1. memref without an affine_layout map being the convention to represent the canonical contiguous buffer of size equal to "the product of sizes"
2. memref with an identity affine_map layout being by convention equivalent to no affine_layout
3. memref with a strided layout that are currently represented by an affine_map.

This also intersects with the fact that a memref layout is provisioned to allow an arbitrarily complex composition of affine maps that I won't even get into.

I forget between the equivalence between 1. and 2. predates the strided memref or not (I think so but I am not 100% sure).

Things will get simpler once strided memref get back to using a different layout attribute offset:x, strides:[x,x,x] as originally intended.
However there are still many skeletons left because in the absence of a richer type/attribute system, it is unknown how the x in the strides relate to the ? in the sizes.
Properly encoding that relationship is key to being airtight in representing contiguous / non-contiguous layouts and reshapes with dynamic shapes.

Based on this info, please feel free to update docs so that it is less surprising to your future self and to others who will approach from the same angle as you did.

Thanks much!

Cant see how this is related to the change here, but I did run into this case today. The below operation currently fails roundtripping

%1 = memref.subview %0 [] [] [] : memref<?x?xf32> to memref<?x?xf32>

since the result subview expects a strided a layout when it shouldnt. I think it is because the subview creation method does not actually adhere to 1 and 2. I was planning to fix that.

[Benoit]

rebased

[Benoit]

I don't have permissions to submit this myself.

[Benoit]

rebased and confirmed it is still needed to use vector transfer flattening.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

public:
   using ConvertOpToLLVMPattern<ReshapeOp>::ConvertOpToLLVMPattern;
   using ReshapeOpAdaptor = typename ReshapeOp::Adaptor;
 
   LogicalResult
   matchAndRewrite(ReshapeOp reshapeOp, typename ReshapeOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     MemRefType dstType = reshapeOp.getResultType();
     MemRefType srcType = reshapeOp.getSrcType();
+
+    // The condition on the layouts can be ignored when all shapes are static.
+    if (!srcType.hasStaticShape() || !dstType.hasStaticShape()) {
       if (!srcType.getLayout().isIdentity() ||
           !dstType.getLayout().isIdentity()) {
-      return rewriter.notifyMatchFailure(reshapeOp,
-                                         "only empty layout map is supported");
+        return rewriter.notifyMatchFailure(
+            reshapeOp, "only empty layout map is supported");
+      }
     }
 
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     if (failed(getStridesAndOffset(dstType, strides, offset))) {
       return rewriter.notifyMatchFailure(
           reshapeOp, "failed to get stride and offset exprs");
     }

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

func @atomic_rmw(%I : memref<10xi32>, %ival : i32, %F : memref<10xf32>, %fval : f32, %i : index) {
   memref.atomic_rmw addf %fval, %F[%i] : (f32, memref<10xf32>) -> f32
   // CHECK: llvm.atomicrmw fadd %{{.*}}, %{{.*}} acq_rel
   memref.atomic_rmw ori %ival, %I[%i] : (i32, memref<10xi32>) -> i32
   // CHECK: llvm.atomicrmw _or %{{.*}}, %{{.*}} acq_rel
   memref.atomic_rmw andi %ival, %I[%i] : (i32, memref<10xi32>) -> i32
   // CHECK: llvm.atomicrmw _and %{{.*}}, %{{.*}} acq_rel
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @collapse_static_shape_with_non_identity_layout
+func @collapse_static_shape_with_non_identity_layout(%arg: memref<1x1x8x8xf32, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 64 + s0 + d1 * 64 + d2 * 8 + d3)>>) -> memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>> {
+// CHECK-NOT: memref.collapse_shape
+  %1 = memref.collapse_shape %arg [[0, 1, 2, 3]] : memref<1x1x8x8xf32, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 64 + s0 + d1 * 64 + d2 * 8 + d3)>> into memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+  return %1 : memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+}