The final update of static checker for linalg operations with the maximum and minimum indices

[• inho9606]

Summary

This version takes the real maximum and minimum indices of each affine expressions binding to each dims. So, we can check if the index ranges are valid or not.

[• inho9606]

This version can check if the index ranges are in bound of the sizes of dimensions using getStaticMaximumIndex and getStaticMinimumIndex methods. The methods get the appropriate dims and substitude them to Affine Expressions using AffineMap.compose method. For example, if we have (-d0 + d1) and d0 and d1 are from 0 to 1, then the maximum index method has 0 for d0 and 1 for d1 because the negative values get larger as its absolute values are smaller. And, the minimum index method has 1 for d0 and 0 for d1.
This version uses AffineMap.compose to get each affine expressions' constant value, but the map.compose method returns all results of expressions of the map. I tried to make a new compose method returning only the result of the given expression, but it seemed I still need AffineMap.. Is there any way to get the only target expression's value? I tried to use getAffineConstantExpr method, but it did not return the exact result..

[hanchung]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

94

We usually pass SmallVectorImpl as function argument.

98–109

I don't get the idea easily, maybe need some more comments? Do you have any assumption here?

[• inho9606]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

94

Ah, like this?
SmallVectorImpl<SmallVector<int64_t, 4>>?

98–109

My idea is something like this:
The goal is to get the maximum index and minimum index in the range (result of the given affine expression). For example, let's say we have affine_map<(d0, d1) -> (d0, d1, -d0 + d1)> binding to memref<3x3x3xf32>.
According to the createLoopRanges method, we can get d0 = 3 and d1 = 3, which is 0 <= d0 < 3 and 0 <= d1 < 3.
To maximize the expression -d0 + d1, we need 0 for d0 and 2 for d1, 0 + 2 = 2. To minimize the expression, we need 2 for d0 and 0 for d1, (-2 + 0 = -2), which is an invalid index.
So, if the affine expression has dims multiplied by a negative constant, then we should substitude appropriate values (0 or range) to the dims in order to get the maximum or minimum value. And I found affine expression writes dims * constant even if the given expression is constant * dim. For example, affine_map<(d0) -> (-d0)> is written (d0 * -1).
AffineExpr.walk method allows me to check the AffineExprKinds of its AffineExpr. So, if it is 'd0 * -1', then I can get its result step by step.
d0 == AffineExprKind::DimID --> d0
-1 == AffineExprKind::constant --> -1
'* == AffineExprKind::mul --> d0 * -1 (Please remove ', it automatically makes lists starting with *..)
So, to get the kind of operator (add or multiply), I had to keep two operands before reaching out the operator using variables named a and b (I know the names are not clear,, could you suggest me better one?). I tried to find dividing cases in tests, but couldn't find them. So it only works with add and multiplication.. Is it fine? Please let me know if I missed something..
Thanks

[hanchung]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

94

I meant SmallVectorImpl<int64_t>.

98–109

I'd like to point out this again. They are rare cases. I don't see any use cases of this before. This increases overheads in verifying operations. However, we don't see any such case, which means that we are just adding overheads in verification. I also suspect we won't see them in the near future. Can we add the check when this is really a case?

Other concerns on my side:

1. The example is very trivial. I'm afraid that we miss many cases.
2. There are other kinds of AffineExpr, like Mod, FloorDiv, etc. What happened in those cases is undefined in your implementation.
3. Some affine expression won't be simplified during parsing, e.g., affine_map<(d0, d1) -> (-(d0 + d1))>. Maybe simplifyAffineExpr would work for this case.

[hkim15]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

520

Logically, en.value().getDimSize() makes sense only when en.value().isDynamicDim(j) is false. If so, I recommend checking en.value().isDynamicDim(j) before calling en.value().getDimSize().

[hkim15]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

520

If shapedDimSize == 0, any access to this tensor is invalid since # of elements in this tensor is 0. I think we should report an error instead of ignoring this and "continue" to the next dimension.

[• inho9606]

1. Updating D102305: The final update of static checker for linalg operations with the maximum and minimum indices #
2. Enter a brief description of the changes included in this update.
3. The first line is used as subject, next lines as comment.

Rebase for sync

[• inho9606]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

98–109

I see. Thanks for letting me the concerns. I just found the case that the startIndices and endIndices might not detect, and wanted to try making the checker as perfect as possible. But my implementation might have other issues as you said.. I thought the parser could simplify all affine expressions, and mod or floordiv might not be used here because I couldn't find such testcases..
Then, can I close this revision after discussing other suggestions with @hkim15?

520

Yes, I think it can still be checked as the invalid one even although we don't check if the shapedDimSize == 0 here because of line number 494 of the original code, or line number 530-531 of the current patch. If the DimSize is 0, then the index would be -1, so the conditions would detect it and generate the error message. But, @hanchung and I found a testcase using 0xf32 tensor, so we had to choose to modify the testcase or ignore the case.

func @dce_zero_memref(%arg0 : memref<0xf32>, %arg1: tensor<0xf32>) -> tensor<0xf32> {

// memref<0x32> is expected to be dce'ed
linalg.copy(%arg0, %arg0): memref<0xf32>, memref<0xf32>

// tensor<0xf32> cannot be dce'ed
%1 = linalg.generic #trait outs(%arg1 : tensor<0xf32>) {
^bb(%0: f32) :
  linalg.yield %0 : f32
} -> tensor<0xf32>

return %1: tensor<0xf32>

}

@hanchung What do you think about this?

[hanchung]

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

98–109

the parser could simplify all affine expressions

Some could be simplified, but some could not. I looked into some tests of AffineMap and realized it.

mod or floordiv might not be used here because I couldn't find such testcases.

They are not well defined I think. We can have them, but we don't recommend it. Actually the case you're trying to check can't be found in the repo right?

can I close this revision after discussing other suggestions with @hkim15

Sure thing, feel free to close it.

520

They were come from https://reviews.llvm.org/D85413

I don't have much context about it, but there are some reasons to having it.

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

static bool isAddMul(Block &block) {
   bool success = un(res, add);
   // One of the operands of add traces back to argC, the other to the mul.
   success |= (un(c1, argC) && un(c2, mul)) || ((un(c1, mul)) && un(c2, argC));
   // One of the operands of mul traces back to argA, the other to argB.
   success |= (un(a, argA) && un(b, argB)) || ((un(a, argB)) && un(b, argA));
   return success;
 }
 
+/// Returns the maximum and minimum index of the given Affine Expression.
+/// This is used for static bound checker.
+static int64_t getStaticMaximumIndex(AffineMap &indexingMap,
+                                     SmallVector<int64_t, 4> loopRanges,
+                                     unsigned idx) {
+  AffineExpr a, b;
+  auto expr = indexingMap.getResult(idx);
+  expr.walk([&](AffineExpr e) {
+    if (auto d = e.dyn_cast<AffineDimExpr>())
+      a = d;
+    else if (auto c = e.dyn_cast<AffineConstantExpr>())
+      b = c;
+    else {
+      if (e.getKind() == AffineExprKind::Mul &&
+          b.dyn_cast<AffineConstantExpr>().getValue() < 0 &&
+          a.dyn_cast<AffineDimExpr>())
+        loopRanges[a.dyn_cast<AffineDimExpr>().getPosition()] = 0;
+    }
+  });
+  return indexingMap.compose(loopRanges)[idx];
+}
+
+static int64_t getStaticMinimumIndex(AffineMap &indexingMap,
+                                     SmallVector<int64_t, 4> &loopRanges,
+                                     unsigned idx) {
+  SmallVector<int64_t, 4> ranges(loopRanges.size(), 0);
+  AffineExpr a, b;
+  auto expr = indexingMap.getResult(idx);
+  expr.walk([&](AffineExpr e) {
+    if (auto d = e.dyn_cast<AffineDimExpr>())
+      a = d;
+    else if (auto c = e.dyn_cast<AffineConstantExpr>())
+      b = c;
+    else {
+      if (e.getKind() == AffineExprKind::Mul &&
+          b.dyn_cast<AffineConstantExpr>().getValue() < 0 &&
+          a.dyn_cast<AffineDimExpr>())
+        ranges[a.dyn_cast<AffineDimExpr>().getPosition()] =
+            loopRanges[a.dyn_cast<AffineDimExpr>().getPosition()];
+    }
+  });
+  return indexingMap.compose(ranges)[idx];
+}
+
 enum MatchContractionResult {
   Success = 0,
   NotLinalgOp,
   WrongNumOperands,
   NoReduction,
   NotProjectedPermutations,
   NotAddMul
 };
if (std::get<0>(it).getElementType() != std::get<1>(it).getType())
       return op->emitError("expected type of bb argument #")
              << (idx + numBBIvs) << " (" << std::get<1>(it).getType() << ")"
              << " to match element type of corresponding shaped operand ("
              << std::get<0>(it).getElementType() << ")";
     ++idx;
   }
 
   // Check if given shapes match to inferred shapes.
-  Optional<SmallVector<int64_t, 4>> endLoopRangeValues =
-      linalgOp.getStaticLoopRanges();
-  if (!endLoopRangeValues)
+  Optional<SmallVector<int64_t, 4>> loopRanges = linalgOp.getStaticLoopRanges();
+  if (!loopRanges)
     return linalgOp.emitError("unable to find loop range for operation");
-  SmallVector<int64_t, 4> startLoopRangeValues((*endLoopRangeValues).size(), 0);
 
   // Verify only static cases since we can't get exact dimension sizes and loop
   // ranges for dynamic cases in this stage.
-  if (llvm::none_of(*endLoopRangeValues, [](int64_t &range) {
+  if (llvm::none_of(*loopRanges, [](int64_t &range) {
         return range == ShapedType::kDynamicSize;
       })) {
-    for (int64_t &range : *endLoopRangeValues)
+    for (int64_t &range : *loopRanges)
       range -= 1;
     for (const auto &en : llvm::enumerate(linalgOp.getShapedOperandTypes())) {
-      auto startIndices =
-          indexingMaps[en.index()].compose(startLoopRangeValues);
-      auto endIndices = indexingMaps[en.index()].compose(*endLoopRangeValues);
       for (auto j : llvm::seq<unsigned>(0, en.value().getRank())) {
 
         // Ignore dynamic dimension or the case that the dimension size is 0
         auto shapedDimSize = en.value().getDimSize(j);
         if (en.value().isDynamicDim(j) || shapedDimSize == 0)
           continue;
 
-        // The first index or last index should be the maximum or the minimum in
-        // the inferred index ranges since the range is increasing or
-        // decreasing. The size of dimensions of shaped operands and the maximum
-        // value + 1 in the inferred range should be the same. But, for now we
+        // The size of dimensions of shaped operands and the maximum
+        // index + 1 in the inferred range should be the same. But, for now we
         // check if the inferred ranges are in boundary of shaped operands' size
         // or not in case that Affine Expressions are complicated such as d0 * 3
         // + d1 since it is not easy to handle the issues.
-        // Found the case that this solution can't check, for example, (d0, d1)
-        // -> (d1 - d0)
-        auto inferredDimSize = std::max(startIndices[j], endIndices[j]) + 1;
-        if (std::min(startIndices[j], endIndices[j]) < 0) {
+        auto inferredDimSize =
+            getStaticMaximumIndex(indexingMaps[en.index()], *loopRanges, j) + 1;
+        if (getStaticMinimumIndex(indexingMaps[en.index()], *loopRanges, j) <
+            0) {
           std::string mapStr;
           {
             llvm::raw_string_ostream os(mapStr);
             os << indexingMaps[en.index()];
           }
           return linalgOp.emitError(
                      "unexpected result less than 0 at expression #")
                  << j << " in " << mapStr;

mlir/test/Dialect/Linalg/invalid.mlir

 func @invalid_reverse(%A: memref<5xf32>, %B: memref<5xf32>) {
   // expected-error @+1 {{unexpected result less than 0 at expression #0 in}}
   linalg.generic #attrs ins(%A: memref<5xf32>) outs(%B: memref<5xf32>) {
     ^bb0(%a: f32, %b: f32):
     linalg.yield %a : f32
   } 
   return
 }
+
+// -----
+
+#attrs = {
+  indexing_maps = [
+    affine_map<(i, j) -> (-i + j)>,
+    affine_map<(i, j) -> (j)>,
+    affine_map<(i, j) -> (i)>
+  ],
+  iterator_types = ["parallel", "parallel"]
+}
+
+func @accessing_index_less_than_0(%A: memref<2xf32>, %B: memref<2xf32>, %C: memref<2xf32>) {
+  // expected-error @+1 {{unexpected result less than 0 at expression #0 in}}
+  linalg.generic #attrs ins(%A, %B: memref<2xf32>, memref<2xf32>) outs(%C: memref<2xf32>) {
+    ^bb0(%a: f32, %b: f32, %c: f32):
+    linalg.yield %b : f32
+  } 
+  return
+}