[mlir][Linalg] Reimplement and extend getStridesAndOffset

[nicolasvasilache]

Summary

This diff reimplements getStridesAndOffset in a significantly simpler way by operating on the AffineExpr and calling into simplifyAffineExpr instead of rolling its own saturating arithmetic.

As a consequence it becomes quite simple to extend the behavior of getStridesAndOffset to encompass more cases by manipulating the AffineExpr more directly.
The divisions are still filtered out and continue to yield fully dynamic strides.
Simplifying the divisions is left for a later time if compelling use cases arise.

Relevant tests are added.

[nicolasvasilache]

Add calls to simplifyAffineExpr.

[merge_guards_bot]

Unit tests: pass. 61127 tests passed, 0 failed and 728 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61175 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

Do you have an estimate on the impact of doing this? If I recall correctly, simplifyAffineExpr is kind of expensive/heavy.

[nicolasvasilache]

@rriddle I am not concerned because the heavy cases involve the introduction of local variables that come from divisions/modulo and FourierMotzkin.
Concretely I do not see a measurable effect when running the tests.

There are 2 additional orthogonal aspects:

1. I proposed in the past that we have a very strict type layering of AffineExpr to prevent people from accidentally pulling in expensive computations silently. This was rejected at the time but is probably due for a refresh now.
2. Since we're uniquing properly, at some point it will make sense to cache the a map<AffineExpr, simplied AffineExpr>, I had done some measurements more than a year ago and a good 75% of invocations were redundant.

These are low-pri for me to the overall trajectory but if they become problematic they should of course be addressed.

[ftynse]

mlir/lib/IR/StandardTypes.cpp

460

This looks like it could accept a hitherto non-canonical form of offset with multiple symbols -- N+M+K+42 -- and treat it as dynamic offset. Is this intentional? If so, please document. Otherwise, there may be a check on seenOffset missing somewhere.

472

If this dim hasn't been seen yet, can strides[dim] be something other than zero? If not, just assign 1 to it.

556–557

(maybe later) I'd consider extracting a utility function from makeCanonicalStridedLayoutExpr that gives you strides as affine expressions, and reusing it here for immediate return instead of doing rounds of simplification on the expression you already know constructed from strides.

559

There's already a check in line 530

596

There's a check that it has exactly one element in line 530

600

There's a check above already

[nicolasvasilache]

Simplifying further by dropping seen and failed.
Addressing review comments.

[nicolasvasilache]

mlir/lib/IR/StandardTypes.cpp

460

Yes it is intentional: previously as I was doing my own saturation arithmetic I needed to special case.
Now I just compute the symbolic expr and simplify in the end.

472

Now that I accumulate AffineExpr symbolically, the assertion on seen is incorrect, thanks!
In fact I can remove seen everywhere and just check against 0 in the end.

556–557

This was already almost the case.
I pushed simplification a bit further and got to what you suggested.

596

It can be empty.

[merge_guards_bot]

Unit tests: pass. 61175 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

mlir/lib/IR/StandardTypes.cpp

460

Then I wouildn't call this diff "NFC", and rather add tests and user-visible documentation for these new cases. IMO, it deserves mentioning that (N+M+K + N*i0 + M*i1 - K*i1) gets stored as three values that are _not_ N,M,K.

465

This never fails, maybe just keep void as return type?

573

Nit: I remember writing code that drops identity maps from the layout, if it's still there, we shouldn't see them here. I'm fine with this code in any case.

[rriddle]

mlir/lib/IR/StandardTypes.cpp

605

failed

[nicolasvasilache]

Address review comments, extend the behavior and add relevant tests.

mlir/lib/IR/StandardTypes.cpp

460

Ok then .. benhavior extend and tests added.

573

ah cool, I'll update to assertions everywhere.

[nicolasvasilache]

Delete spurious commented out code and include.

[merge_guards_bot]

Unit tests: pass. 61236 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[merge_guards_bot]

Unit tests: pass. 61236 tests passed, 0 failed and 729 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[ftynse]

Thanks!

[nicolasvasilache]

+@mravishankar

mlir/lib/IR/StandardTypes.cpp

 //===- StandardTypes.cpp - MLIR Standard Type Classes ---------------------===//
 //
 // Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/StandardTypes.h"
 #include "TypeDetail.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/STLExtras.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 // Type
 //===----------------------------------------------------------------------===//
 
if (!dynamicPoisonBit) {
         // From now on bail into dynamic mode.
         dynamicPoisonBit = true;
       continue;
     }
     // Dynamic case, new symbol for each new stride.
     auto sym = getAffineSymbolExpr(nSymbols++, context);
     expr = expr ? expr + d * sym : d * sym;
   }
-  return expr;
+  return simplifyAffineExpr(expr, rank, nSymbols);
 }
 
-// Factored out common logic to update `strides` and `seen` for `dim` with value
-// `val`. This handles both saturated and unsaturated cases.
-static void accumulateStrides(MutableArrayRef<int64_t> strides,
-                              MutableArrayRef<bool> seen, unsigned pos,
-                              int64_t val) {
-  if (!seen[pos]) {
-    // Newly seen case, sets value
-    strides[pos] = val;
-    seen[pos] = true;
-    return;
-  }
-  if (strides[pos] != MemRefType::getDynamicStrideOrOffset())
-    // Already seen case accumulates unless they are already saturated.
-    strides[pos] += val;
-}
-
-// This sums multiple offsets as they are seen. In the particular case of
-// accumulating a dynamic offset with either a static of dynamic one, this
-// saturates to MemRefType::getDynamicStrideOrOffset().
-static void accumulateOffset(int64_t &offset, bool &seenOffset, int64_t val) {
-  if (!seenOffset) {
-    // Newly seen case, sets value
-    offset = val;
-    seenOffset = true;
-    return;
-  }
-  if (offset != MemRefType::getDynamicStrideOrOffset())
-    // Already seen case accumulates unless they are already saturated.
-    offset += val;
-}
-
-/// Takes a single AffineExpr `e` and populates the `strides` and `seen` arrays
-/// with the strides values for each dim position and whether a value exists at
-/// that position, respectively.
+// Fallback cases for terminal dim/sym/cst that are not part of a binary op (
+// i.e. single term). Accumulate the AffineExpr into the existing one.
+static void extractStridesFromTerm(AffineExpr e,
+                                   AffineExpr multiplicativeFactor,
+                                   MutableArrayRef<AffineExpr> strides,
+                                   AffineExpr &offset) {
+  // llvm::errs() << "\nFac: " << multiplicativeFactor;
+  if (auto dim = e.dyn_cast<AffineDimExpr>())
+    strides[dim.getPosition()] =
+        strides[dim.getPosition()] + multiplicativeFactor;
+  else
+    offset = offset + e * multiplicativeFactor;
+}
+
+/// Takes a single AffineExpr `e` and populates the `strides` array with the
+/// strides expressions for each dim position.
 /// The convention is that the strides for dimensions d0, .. dn appear in
 /// order to make indexing intuitive into the result.
-static void extractStrides(AffineExpr e, MutableArrayRef<int64_t> strides,
-                           int64_t &offset, MutableArrayRef<bool> seen,
-                           bool &seenOffset, bool &failed) {
+static LogicalResult extractStrides(AffineExpr e,
+                                    AffineExpr multiplicativeFactor,
+                                    MutableArrayRef<AffineExpr> strides,
+                                    AffineExpr &offset) {
+  // e.print(llvm::errs() << "\nIterate: ");
   auto bin = e.dyn_cast<AffineBinaryOpExpr>();
-  if (!bin)
-    return;
+  if (!bin) {
+    extractStridesFromTerm(e, multiplicativeFactor, strides, offset);
+    return success();
+  }
 
   if (bin.getKind() == AffineExprKind::CeilDiv ||
       bin.getKind() == AffineExprKind::FloorDiv ||
-      bin.getKind() == AffineExprKind::Mod) {
-    failed = true;
-    return;
-  }
+      bin.getKind() == AffineExprKind::Mod)
+    return failure();
+
   if (bin.getKind() == AffineExprKind::Mul) {
-    // LHS may be more complex than just a single dim (e.g. multiple syms and
-    // dims). Bail out for now and revisit when we have evidence this is needed.
     auto dim = bin.getLHS().dyn_cast<AffineDimExpr>();
-    if (!dim) {
-      failed = true;
-      return;
-    }
-    auto cst = bin.getRHS().dyn_cast<AffineConstantExpr>();
-    if (!cst) {
-      strides[dim.getPosition()] = MemRefType::getDynamicStrideOrOffset();
-      seen[dim.getPosition()] = true;
-    } else {
-      accumulateStrides(strides, seen, dim.getPosition(), cst.getValue());
+    if (dim) {
+      strides[dim.getPosition()] =
+          strides[dim.getPosition()] + bin.getRHS() * multiplicativeFactor;
+      return success();
     }
-    return;
+    // LHS and RHS may both contain complex expressions of dims. Try one path
+    // and if it fails try the other. This is guaranteed to succeed because
+    // only one path may have a `dim`, otherwise this is not an AffineExpr in
+    // the first place.
+    // bin.getLHS().print(llvm::errs() << "\nLHS: ");
+    // bin.getRHS().print(llvm::errs() << "\nRHS: ");
+    if (bin.getLHS().isSymbolicOrConstant())
+      return extractStrides(bin.getRHS(), multiplicativeFactor * bin.getLHS(),
+                            strides, offset);
+    return extractStrides(bin.getLHS(), multiplicativeFactor * bin.getRHS(),
+                          strides, offset);
   }
+
   if (bin.getKind() == AffineExprKind::Add) {
-    for (auto e : {bin.getLHS(), bin.getRHS()}) {
-      if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
-        // Independent constants cumulate.
-        accumulateOffset(offset, seenOffset, cst.getValue());
-      } else if (auto sym = e.dyn_cast<AffineSymbolExpr>()) {
-        // Independent symbols saturate.
-        offset = MemRefType::getDynamicStrideOrOffset();
-        seenOffset = true;
-      } else if (auto dim = e.dyn_cast<AffineDimExpr>()) {
-        // Independent symbols cumulate 1.
-        accumulateStrides(strides, seen, dim.getPosition(), 1);
-      }
-      // Sum of binary ops dispatch to the respective exprs.
-    }
-    return;
-  }
-  llvm_unreachable("unexpected binary operation");
+    auto res1 =
+        extractStrides(bin.getLHS(), multiplicativeFactor, strides, offset);
+    auto res2 =
+        extractStrides(bin.getRHS(), multiplicativeFactor, strides, offset);
+    return success(succeeded(res1) && succeeded(res2));
   }
 
-// Fallback cases for terminal dim/sym/cst that are not part of a binary op (
-// i.e. single term).
-static void extractStridesFromTerm(AffineExpr e,
-                                   MutableArrayRef<int64_t> strides,
-                                   int64_t &offset, MutableArrayRef<bool> seen,
-                                   bool &seenOffset) {
-  if (auto cst = e.dyn_cast<AffineConstantExpr>()) {
-    assert(!seenOffset && "unexpected `seen` bit with single term");
-    offset = cst.getValue();
-    seenOffset = true;
-    return;
-  }
-  if (auto sym = e.dyn_cast<AffineSymbolExpr>()) {
-    assert(!seenOffset && "unexpected `seen` bit with single term");
-    offset = MemRefType::getDynamicStrideOrOffset();
-    seenOffset = true;
-    return;
-  }
-  if (auto dim = e.dyn_cast<AffineDimExpr>()) {
-    assert(!seen[dim.getPosition()] &&
-           "unexpected `seen` bit with single term");
-    strides[dim.getPosition()] = 1;
-    seen[dim.getPosition()] = true;
-    return;
-  }
   llvm_unreachable("unexpected binary operation");
 }
 
-LogicalResult mlir::getStridesAndOffset(MemRefType t,
-                                        SmallVectorImpl<int64_t> &strides,
-                                        int64_t &offset) {
+static LogicalResult getStridesAndOffset(MemRefType t,
+                                         SmallVectorImpl<AffineExpr> &strides,
+                                         AffineExpr &offset) {
   auto affineMaps = t.getAffineMaps();
   // For now strides are only computed on a single affine map with a single
   // result (i.e. the closed subset of linearization maps that are compatible
   // with striding semantics).
   // TODO(ntv): support more forms on a per-need basis.
   if (affineMaps.size() > 1)
     return failure();
-  AffineExpr stridedExpr;
-  if (affineMaps.empty() || affineMaps[0].isIdentity()) {
-    if (t.getRank() == 0) {
-      // Handle 0-D corner case.
-      offset = 0;
-      return success();
-    }
-    stridedExpr = makeCanonicalStridedLayoutExpr(t.getShape(), t.getContext());
-  } else if (affineMaps[0].getNumResults() == 1) {
-    stridedExpr = affineMaps[0].getResult(0);
-  }
-  if (!stridedExpr)
+  if (affineMaps.size() == 1 && affineMaps[0].getNumResults() != 1)
     return failure();
 
-  bool failed = false;
-  strides = SmallVector<int64_t, 4>(t.getRank(), 0);
-  bool seenOffset = false;
-  SmallVector<bool, 4> seen(t.getRank(), false);
-  if (stridedExpr.isa<AffineBinaryOpExpr>()) {
-    stridedExpr.walk([&](AffineExpr e) {
-      if (!failed)
-        extractStrides(e, strides, offset, seen, seenOffset, failed);
-    });
-  } else {
-    extractStridesFromTerm(stridedExpr, strides, offset, seen, seenOffset);
+  auto zero = getAffineConstantExpr(0, t.getContext());
+  auto one = getAffineConstantExpr(1, t.getContext());
+  offset = zero;
+  strides.assign(t.getRank(), zero);
+
+  AffineMap m;
+  if (!affineMaps.empty()) {
+    m = affineMaps.front();
+    assert(!m.isIdentity() && "unexpected identity map");
+  }
+
+  // Canonical case for empty map.
+  if (!m) {
+    // 0-D corner case, offset is already 0.
+    if (t.getRank() == 0)
+      return success();
+    auto stridedExpr =
+        makeCanonicalStridedLayoutExpr(t.getShape(), t.getContext());
+    if (succeeded(extractStrides(stridedExpr, one, strides, offset)))
+      return success();
+    assert(false && "unexpected failure: extract strides in canonical layout");
   }
 
-  // Constant offset may not be present in `stridedExpr` which means it is
-  // implicitly 0.
-  if (!seenOffset)
-    offset = 0;
+  // Non-canonical case requires more work.
+  auto stridedExpr =
+      simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
+  // stridedExpr.print(llvm::errs() << "\nOrig expr: ");
+  if (failed(extractStrides(stridedExpr, one, strides, offset))) {
+    offset = AffineExpr();
+    strides.clear();
+    return failure();
+  }
 
-  if (failed || !llvm::all_of(seen, [](bool b) { return b; })) {
+  // Simplify results to allow folding to constants and simple checks.
+  unsigned numDims = m.getNumDims();
+  unsigned numSymbols = m.getNumSymbols();
+  offset = simplifyAffineExpr(offset, numDims, numSymbols);
+  for (auto &stride : strides) {
+    // stride.print(llvm::errs() << "\nStride: ");
+    stride = simplifyAffineExpr(stride, numDims, numSymbols);
+  }
+
+  /// In practice, a strided memref must be internally non-aliasing. Test
+  /// against 0 as a proxy.
+  /// TODO(ntv) static cases can have more advanced checks.
+  /// TODO(ntv) dynamic cases would require a way to compare symbolic
+  /// expressions and would probably need an affine set context propagated
+  /// everywhere.
+  if (llvm::any_of(strides, [](AffineExpr e) {
+        return e == getAffineConstantExpr(0, e.getContext());
+      })) {
+    offset = AffineExpr();
     strides.clear();
     return failure();
   }
 
   return success();
 }
 
+LogicalResult mlir::getStridesAndOffset(MemRefType t,
+                                        SmallVectorImpl<int64_t> &strides,
+                                        int64_t &offset) {
+  AffineExpr offsetExpr;
+  SmallVector<AffineExpr, 4> strideExprs;
+  if (failed(::getStridesAndOffset(t, strideExprs, offsetExpr)))
+    return failure();
+  if (auto cst = offsetExpr.dyn_cast<AffineConstantExpr>())
+    offset = cst.getValue();
+  else
+    offset = ShapedType::kDynamicStrideOrOffset;
+  for (auto e : strideExprs) {
+    if (auto c = e.dyn_cast<AffineConstantExpr>())
+      strides.push_back(c.getValue());
+    else
+      strides.push_back(ShapedType::kDynamicStrideOrOffset);
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 /// ComplexType
 //===----------------------------------------------------------------------===//
 
 ComplexType ComplexType::get(Type elementType) {
   return Base::get(elementType.getContext(), StandardTypes::Complex,
                    elementType);
 }

mlir/test/AffineOps/memref-stride-calculation.mlir

// CHECK: MemRefType offset: 123 strides:
   %100 = alloc(%0, %0)[%0, %0] : memref<?x?x16xf32, (i, j, k)[M, N]->(i + j, j, k), (i, j, k)[M, N]->(M * i + N * j + k + 1)>
 // CHECK: MemRefType memref<?x?x16xf32, (d0, d1, d2)[s0, s1] -> (d0 + d1, d1, d2), (d0, d1, d2)[s0, s1] -> (d0 * s0 + d1 * s1 + d2 + 1)> cannot be converted to strided form
   %101 = alloc() : memref<3x4x5xf32, (i, j, k)->(i floordiv 4 + j + k)>
 // CHECK: MemRefType memref<3x4x5xf32, (d0, d1, d2) -> (d0 floordiv 4 + d1 + d2)> cannot be converted to strided form
   %102 = alloc() : memref<3x4x5xf32, (i, j, k)->(i ceildiv 4 + j + k)>
 // CHECK: MemRefType memref<3x4x5xf32, (d0, d1, d2) -> (d0 ceildiv 4 + d1 + d2)> cannot be converted to strided form
   %103 = alloc() : memref<3x4x5xf32, (i, j, k)->(i mod 4 + j + k)>
 // CHECK: MemRefType memref<3x4x5xf32, (d0, d1, d2) -> (d0 mod 4 + d1 + d2)> cannot be converted to strided form
+
+  %200 = alloc()[%0, %0, %0] : memref<3x4x5xf32, (i, j, k)[M, N, K]->(M * i + N * i + N * j + K * k - (M + N - 20)* i)>
+  // CHECK: MemRefType offset: 0 strides: 20, ?, ?
+  %201 = alloc()[%0, %0, %0] : memref<3x4x5xf32, (i, j, k)[M, N, K]->(M * i + N * i + N * K * j + K * K * k - (M + N - 20) * (i + 1))>
+  // CHECK: MemRefType offset: ? strides: 20, ?, ?
+  %202 = alloc()[%0, %0, %0] : memref<3x4x5xf32, (i, j, k)[M, N, K]->(M * (i + 1) + j + k - M)>
+  // CHECK: MemRefType offset: 0 strides: ?, 1, 1
+  %203 = alloc()[%0, %0, %0] : memref<3x4x5xf32, (i, j, k)[M, N, K]->(M + M * (i + N * (j + K * k)))>
+  // CHECK: MemRefType offset: ? strides: ?, ?, ?
+
   return
 }