[mlir][Toy][LowerToAffineLoops] Handle tensors of rank 0

[djtodoro]

Summary

This will fix the case:

toyc -emit=jit test.toy

def main() {
  var a = 1;
  print(a);
}

Without this patch it would trigger an assertion.

[bondhugula]

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

159

Nit: !valueShape.empty() if you prefer.

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

159

Likewise.

[mehdi_amini]

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

159

+1 for empty() :)

[djtodoro]

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

159

Makes sense to me. Thanks :)

[djtodoro]

-addressing comments

[mehdi_amini]

Still LGTM, let me know if you need someone to land it for you

[mehdi_amini]

(Seems like you didn't use arc, so I'd need the info on what to put in the author field of the commit)

[djtodoro]

I have the access, thanks. :)
I’ve just started learning MLIR and come up with this, but I actually work on some other LLVM parts :)

[mehdi_amini]

Ah great, thanks for extending your contributions to MLIR then :)

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

LogicalResult matchAndRewrite(toy::ConstantOp op,
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
+
+    if (valueShape.size()) {
       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,
     // generating a store when the recursion hits the base case.
     SmallVector<Value, 2> indices;
     auto valueIt = constantValue.getValues<FloatAttr>().begin();
     std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {

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

LogicalResult matchAndRewrite(toy::ConstantOp op,
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
+
+    if (valueShape.size()) {
       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,
     // generating a store when the recursion hits the base case.
     SmallVector<Value, 2> indices;
     auto valueIt = constantValue.getValues<FloatAttr>().begin();
     std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {
       // The last dimension is the base case of the recursion, at this point

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

LogicalResult matchAndRewrite(toy::ConstantOp op,
     auto memRefType = convertTensorToMemRef(tensorType);
     auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 
     // We will be generating constant indices up-to the largest dimension.
     // Create these constants up-front to avoid large amounts of redundant
     // operations.
     auto valueShape = memRefType.getShape();
     SmallVector<Value, 8> constantIndices;
+
+    if (valueShape.size()) {
       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,
     // generating a store when the recursion hits the base case.
     SmallVector<Value, 2> indices;
     auto valueIt = constantValue.getValues<FloatAttr>().begin();
     std::function<void(uint64_t)> storeElements = [&](uint64_t dimension) {