[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/include/mlir/Dialect/Affine/Utils.h

 /// and updates all its indexing uses. Returns failure if any of its uses
 /// escape (while leaving the IR in a valid state).
 LogicalResult normalizeMemRef(memref::AllocOp *op);
 
 /// Normalizes `memrefType` so that the affine layout map of the memref is
 /// transformed to an identity map with a new shape being computed for the
 /// normalized memref type and returns it. The old memref type is simplify
 /// returned if the normalization failed.
-MemRefType normalizeMemRefType(MemRefType memrefType,
-                               unsigned numSymbolicOperands);
+MemRefType normalizeMemRefType(MemRefType memrefType);
 
 /// Given an operation, inserts one or more single result affine apply
 /// operations, results of which are exclusively used by this operation.
 /// The operands of these newly created affine apply ops are
 /// guaranteed to be loop iterators or terminal symbols of a function.
 ///
 /// Before
 ///

mlir/lib/Dialect/Affine/Utils/Utils.cpp

 
 // TODO: Currently works for static memrefs with a single layout map.
 LogicalResult mlir::affine::normalizeMemRef(memref::AllocOp *allocOp) {
   MemRefType memrefType = allocOp->getType();
   OpBuilder b(*allocOp);
 
   // Fetch a new memref type after normalizing the old memref to have an
   // identity map layout.
-  MemRefType newMemRefType =
-      normalizeMemRefType(memrefType, allocOp->getSymbolOperands().size());
+  MemRefType newMemRefType = normalizeMemRefType(memrefType);
   if (newMemRefType == memrefType)
     // Either memrefType already had an identity map or the map couldn't be
     // transformed to an identity map.
     return failure();
 
   Value oldMemRef = allocOp->getResult();
 
   SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
LogicalResult mlir::affine::normalizeMemRef(memref::AllocOp *allocOp) {
   assert(llvm::all_of(oldMemRef.getUsers(), [&](Operation *op) {
     return hasSingleEffect<MemoryEffects::Free>(op, oldMemRef);
   }));
   oldMemRef.replaceAllUsesWith(newAlloc);
   allocOp->erase();
   return success();
 }
 
-MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType,
-                                             unsigned numSymbolicOperands) {
+MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
   unsigned rank = memrefType.getRank();
   if (rank == 0)
     return memrefType;
 
   if (memrefType.getLayout().isIdentity()) {
     // Either no maps is associated with this memref or this memref has
     // a trivial (identity) map.
     return memrefType;
   }
   AffineMap layoutMap = memrefType.getLayout().getAffineMap();
+  unsigned numSymbolicOperands = layoutMap.getNumSymbols();
 
   // We don't do any checks for one-to-one'ness; we assume that it is
   // one-to-one.
 
   // Normalize only static memrefs and dynamic memrefs with a tiled-layout map
   // for now.
   // TODO: Normalize the other types of dynamic memrefs.
   SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;

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

for (unsigned argIndex :
     // Check whether argument is of MemRef type. Any other argument type can
     // simply be part of the final function signature.
     if (!memrefType) {
       inputTypes.push_back(argType);
       continue;
     }
     // Fetch a new memref type after normalizing the old memref to have an
     // identity map layout.
-    MemRefType newMemRefType = normalizeMemRefType(memrefType,
-                                                   /*numSymbolicOperands=*/0);
+    MemRefType newMemRefType = normalizeMemRefType(memrefType);
     if (newMemRefType == memrefType || funcOp.isExternal()) {
       // Either memrefType already had an identity map or the map couldn't be
       // transformed to an identity map.
       inputTypes.push_back(newMemRefType);
       continue;
     }
 
     // Insert a new temporary argument with the new memref type.
for (unsigned resIndex :
       // Check whether result is of MemRef type. Any other argument type can
       // simply be part of the final function signature.
       if (!memrefType) {
         resultTypes.push_back(resType);
         continue;
       }
       // Computing a new memref type after normalizing the old memref to have an
       // identity map layout.
-      MemRefType newMemRefType = normalizeMemRefType(memrefType,
-                                                     /*numSymbolicOperands=*/0);
+      MemRefType newMemRefType = normalizeMemRefType(memrefType);
       resultTypes.push_back(newMemRefType);
     }
 
     FunctionType newFuncType =
         FunctionType::get(&getContext(), /*inputs=*/inputTypes,
                           /*results=*/resultTypes);
     // Setting the new function signature for this external function.
     funcOp.setType(newFuncType);
Operation *NormalizeMemRefs::createOpResultsNormalized(func::FuncOp funcOp,
   for (unsigned resIndex : llvm::seq<unsigned>(0, oldOp->getNumResults())) {
     auto resultType = oldOp->getResult(resIndex).getType();
     MemRefType memrefType = dyn_cast<MemRefType>(resultType);
     // Check whether the operation result is MemRef type.
     if (!memrefType) {
       resultTypes.push_back(resultType);
       continue;
     }
+
     // Fetch a new memref type after normalizing the old memref.
-    MemRefType newMemRefType = normalizeMemRefType(memrefType,
-                                                   /*numSymbolicOperands=*/0);
+    MemRefType newMemRefType = normalizeMemRefType(memrefType);
     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: ?>>
+  // CHECK: %{{.*}} = memref.cast %{{.*}} : memref<?x512xf32, strided<[?, ?], offset: ?>> to memref<16x512xf32, 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>
+}