[Mlir-commits] [mlir] da5fe23 - [mlir][LowerToAffineLoops] Handle tensors of rank 0

[Djordje Todorovic]

Author: Djordje Todorovic
Date: 2020-04-06T14:51:03+02:00
New Revision: da5fe23e84c829243162ea2c51fd62d229b19ea4

URL: https://github.com/llvm/llvm-project/commit/da5fe23e84c829243162ea2c51fd62d229b19ea4
DIFF: https://github.com/llvm/llvm-project/commit/da5fe23e84c829243162ea2c51fd62d229b19ea4.diff

LOG: [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

Added: 
    

Modified: 
    mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
    mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp
    mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp b/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
index 249f17b0fc71..0614f3ac043b 100644
--- a/mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp
+++ b/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

diff  --git a/mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp b/mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp
index 249f17b0fc71..4292d14ec3ed 100644
--- a/mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp
+++ b/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,

diff  --git a/mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp b/mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp
index 249f17b0fc71..0614f3ac043b 100644
--- a/mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp
+++ b/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