[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

472

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

476

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.

542

(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.

545

There's already a check in line 530

560

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

564

There's a check above already

[nicolasvasilache]

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

[nicolasvasilache]

mlir/lib/IR/StandardTypes.cpp

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.

476

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.

542

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

560

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

476

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.

481

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

562

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

578

failed

[nicolasvasilache]

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

mlir/lib/IR/StandardTypes.cpp

476

Ok then .. benhavior extend and tests added.

562

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"
 
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;
+// 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<AffineExpr> strides,
+                                   AffineExpr &offset,
+                                   MutableArrayRef<bool> seen,
+                                   bool &seenOffset) {
+  if (auto dim = e.dyn_cast<AffineDimExpr>()) {
+    assert(!seen[dim.getPosition()] &&
+           "unexpected `seen` bit with single term");
+    strides[dim.getPosition()] = strides[dim.getPosition()] + 1;
+    seen[dim.getPosition()] = 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;
+  offset = offset + e;
 }
 
 /// 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.
 /// 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,
+static void extractStrides(AffineExpr e, MutableArrayRef<AffineExpr> strides,
+                           AffineExpr &offset, MutableArrayRef<bool> seen,
                            bool &seenOffset, bool &failed) {
   auto bin = e.dyn_cast<AffineBinaryOpExpr>();
-  if (!bin)
+  if (!bin) {
+    extractStridesFromTerm(e, strides, offset, seen, seenOffset);
     return;
+  }
 
   if (bin.getKind() == AffineExprKind::CeilDiv ||
       bin.getKind() == AffineExprKind::FloorDiv ||
       bin.getKind() == AffineExprKind::Mod) {
     failed = true;
     return;
   }
   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());
-    }
+    strides[dim.getPosition()] = strides[dim.getPosition()] + bin.getRHS();
     return;
   }
   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.
-    }
+    extractStrides(bin.getLHS(), strides, offset, seen, seenOffset, failed);
+    extractStrides(bin.getRHS(), strides, offset, seen, seenOffset, failed);
     return;
   }
   llvm_unreachable("unexpected binary operation");
 }
 
-// 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();
+  if (affineMaps.size() == 1 && affineMaps[0].getNumResults() != 1)
+    return failure();
+
+  auto zero = getAffineConstantExpr(0, t.getContext());
+  offset = zero;
+  strides.assign(t.getRank(), zero);
+
   AffineExpr stridedExpr;
-  if (affineMaps.empty() || affineMaps[0].isIdentity()) {
-    if (t.getRank() == 0) {
-      // Handle 0-D corner case.
-      offset = 0;
+  if (affineMaps.empty()) {
+    // 0-D corner case, offset is already 0.
+    if (t.getRank() == 0)
       return success();
-    }
     stridedExpr = makeCanonicalStridedLayoutExpr(t.getShape(), t.getContext());
-  } else if (affineMaps[0].getNumResults() == 1) {
-    stridedExpr = affineMaps[0].getResult(0);
+  } else {
+    auto m = affineMaps.front();
+    assert(m.getNumResults() == 1);
+    stridedExpr =
+        simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
   }
-  if (!stridedExpr)
-    return failure();
+  assert(stridedExpr);
 
   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);
-  }
-
-  // Constant offset may not be present in `stridedExpr` which means it is
-  // implicitly 0.
-  if (!seenOffset)
-    offset = 0;
+  extractStrides(stridedExpr, strides, offset, seen, seenOffset, failed);
 
   if (failed || !llvm::all_of(seen, [](bool b) { return b; })) {
     strides.clear();
     return failure();
   }
+  if (!affineMaps.empty()) {
+    auto m = affineMaps.front();
+    unsigned numDims = m.getNumDims();
+    unsigned numSymbols = m.getNumSymbols();
+    assert(m.getNumResults() == 1);
+    offset = simplifyAffineExpr(offset, numDims, numSymbols);
+    for (auto &stride : strides)
+      stride = simplifyAffineExpr(stride, numDims, numSymbols);
+  }
 
   return success();
 }
 
+LogicalResult mlir::getStridesAndOffset(MemRefType t,
+                                        SmallVectorImpl<int64_t> &strides,
+                                        int64_t &offset) {
+  AffineExpr offsetExpr;
+  SmallVector<AffineExpr, 4> strideExprs;
+  if (!succeeded(::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);
 }