[mlir][LowerToAffineLoops] Handle tensors of rank 0

This will fix the case:

  $ toyc -emit=jit test.toy
  $ cat test.toy
  def main() {
    var a = 1;
    print(a);
  }

Without this patch it would trigger an assertion.

Differential Revision: https://reviews.llvm.org/D77464

Author: Djordje Todorovic

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

@@ -155,9 +155,15 @@ struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
-    for (auto i : llvm::seq<int64_t>(
-             0, *std::max_element(valueShape.begin(), valueShape.end())))
-      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
+
+    if (!valueShape.empty()) {
+      for (auto i : llvm::seq<int64_t>(
+              0, *std::max_element(valueShape.begin(), valueShape.end())))
+       constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
+    } else {
+      // This is the case of a tensor of rank 0.
+      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
+    }
 
     // The constant operation represents a multi-dimensional constant, so we
     // will need to generate a store for each of the elements. The following

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

@@ -155,10 +155,15 @@ struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
-    for (auto i : llvm::seq<int64_t>(
-             0, *std::max_element(valueShape.begin(), valueShape.end())))
-      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
 
+    if (!valueShape.empty()) {
+      for (auto i : llvm::seq<int64_t>(
+              0, *std::max_element(valueShape.begin(), valueShape.end())))
+       constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
+    } else {
+      // This is the case of a tensor of rank 0.
+      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
+    }
     // The constant operation represents a multi-dimensional constant, so we
     // will need to generate a store for each of the elements. The following
     // functor recursively walks the dimensions of the constant shape,

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

@@ -155,9 +155,15 @@ struct ConstantOpLowering : public OpRewritePattern<toy::ConstantOp> {
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
-    for (auto i : llvm::seq<int64_t>(
-             0, *std::max_element(valueShape.begin(), valueShape.end())))
-      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
+
+    if (!valueShape.empty()) {
+      for (auto i : llvm::seq<int64_t>(
+              0, *std::max_element(valueShape.begin(), valueShape.end())))
+       constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, i));
+    } else {
+      // This is the case of a tensor of rank 0.
+      constantIndices.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
+    }
 
     // The constant operation represents a multi-dimensional constant, so we
     // will need to generate a store for each of the elements. The following