[mlir][memref] Make result normalization aware of the number symbols

[Lewuathe]

Summary

Memref normalization fails to recognize the non-zero symbols used in the memref type itself with strided, offset information. It causes the crash with the type like memref<128x512xf32, strided<[?, ?], offset: ?>>. The original issue is here. https://github.com/llvm/llvm-project/issues/61345

[Lewuathe]

Apply format.

[Lewuathe]

Check the normalization after cast operation.

[Lewuathe]

@aartbik @rriddle @bondhugula Sorry for bothering you but could you take a look at this one when you get a chance?

[Lewuathe]

@nicolasvasilache @springerm @dcaballe Sorry for bothering you time to time but could you review this one when you get a chance?

[springerm]

Sorry, I'm not familiar with this part of the codebase. If we can just take the number of symbols from the memref type, why does this function even have the numSymbolicOperands parameter?

[Lewuathe]

@springerm I'm not 100% sure, but normalizeMemRefType does not always get the number of symbols from the memref type itself. In the following case, we have to provide the number of symbols from the allocOp so that the caller should be able to determine how to get it.

MemRefType newMemRefType =
    normalizeMemRefType(memrefType, allocOp->getSymbolOperands().size());

https://github.com/llvm/llvm-project/blob/d8562e27e05b90d8957d20444c724293ddf1ba0c/mlir/lib/Dialect/Affine/Utils/Utils.cpp#L1724

[Lewuathe]

Can someone take a look at this one?

Looks like this issue is still present.
https://discourse.llvm.org/t/assertion-failed-when-lowering-simple-tensor-dilect-example/71336/3

[Lewuathe]

Could someone review this patch?

[Lewuathe]

@nicolasvasilache @springerm @dcaballe @aartbik @mehdi_amini @rriddle @bondhugula @antiagainst @stephenneuendorffer
Sorry for pinging time to time. Could someone review this patch when available?

[mehdi_amini]

In D150250#4419484, @Lewuathe wrote:

@springerm I'm not 100% sure, but normalizeMemRefType does not always get the number of symbols from the memref type itself. In the following case, we have to provide the number of symbols from the allocOp so that the caller should be able to determine how to get it.

MemRefType newMemRefType =
    normalizeMemRefType(memrefType, allocOp->getSymbolOperands().size());

https://github.com/llvm/llvm-project/blob/d8562e27e05b90d8957d20444c724293ddf1ba0c/mlir/lib/Dialect/Affine/Utils/Utils.cpp#L1724

It would be nice to get an answer about the intent for this.

I added: assert(allocOp->getSymbolOperands().size() == memrefType.getLayout().getAffineMap().getNumSymbols()); here and didn't hit any test failure.

[Lewuathe]

The function was moved by this patch from transform dialect.

The original patch to create the function is here.

https://github.com/llvm/llvm-project/commit/76d07503f0c69f6632e6d8d4736e2a4cb4055a92

There is still no clue about the decision of the function interface.

https://reviews.llvm.org/D84490

@avarmapml @bondhugula Do you remember why normalizeMemRefType function has a second argument to get the number of symbols while we can extract that information from the given memrefType?

[Lewuathe]

If there is no specific reason, I think we can omit the parameter getting numSymbolicOperands.

[Lewuathe]

Omit the number of symbol operands from normalizeMemRefType.

[mehdi_amini]

LGTM, since we didn't find a justification to keep it. It can always be reintroduced if someone needs this.

[Lewuathe]

@mehdi_amini Thank you so much for reviewing!

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

Operation *NormalizeMemRefs::createOpResultsNormalized(func::FuncOp funcOp,
   for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
     auto resultType = oldOp->getResult(resIndex).getType();
     MemRefType memrefType = resultType.dyn_cast<MemRefType>();
     // Check whether the operation result is MemRef type.
     if (!memrefType) {
       resultTypes.push_back(resultType);
       continue;
     }
+
+    AffineMap layoutMap = memrefType.getLayout().getAffineMap();
     // Fetch a new memref type after normalizing the old memref.
     MemRefType newMemRefType = normalizeMemRefType(memrefType,
-                                                   /*numSymbolicOperands=*/0);
+                                                   /*numSymbolicOperands=*/layoutMap.getNumSymbols());
     if (newMemRefType == memrefType) {
       // Either memrefType already had an identity map or the map couldn't
       // be transformed to an identity map.
       resultTypes.push_back(memrefType);
       continue;
     }
     resultTypes.push_back(newMemRefType);
     resultTypeNormalized = true;

mlir/test/Transforms/normalize-memrefs.mlir

 #neg = affine_map<(d0, d1) -> (d0, d1 - 100)>
 // CHECK-LABEL: func.func @neg_map
 func.func @neg_map() -> memref<2x3xf32, #neg> {
   // This isn't a valid map for normalization.
   // CHECK: memref.alloc() : memref<2x3xf32, #{{.*}}>
   %0 = memref.alloc() : memref<2x3xf32, #neg>
   return %0 : memref<2x3xf32, #neg>
 }
+
+// CHECK-LABEL: func @memref_with_strided_offset
+func.func @memref_with_strided_offset(%arg0: tensor<128x512xf32>, %arg1: index, %arg2: index) -> tensor<16x512xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = bufferization.to_memref %arg0 : memref<128x512xf32, strided<[?, ?], offset: ?>>
+  %subview = memref.subview %0[%arg2, 0] [%arg1, 512] [1, 1] : memref<128x512xf32, strided<[?, ?], offset: ?>> to memref<?x512xf32, strided<[?, ?], offset: ?>>
+  %cast = memref.cast %subview : memref<?x512xf32, strided<[?, ?], offset: ?>> to memref<16x512xf32, strided<[?, ?], offset: ?>>
+  %1 = bufferization.to_tensor %cast : memref<16x512xf32, strided<[?, ?], offset: ?>>
+  return %1 : tensor<16x512xf32>
+}