[mlir][linalg] Fixed issue generating reassociation map with Rank-0 types

[rsuderman]

Summary

Rank-0 case causes a graph during linalg reshape operation.

[mravishankar]

Is it possible to get a small repro for the issue. Basically looking for the example where you hit this. I know why this might have happened, but want to make sure it is the expected case.

[rsuderman]

Updated to add tensor reshape canonicalization test.

[rsuderman]

In D102282#2752291, @mravishankar wrote:

Is it possible to get a small repro for the issue. Basically looking for the example where you hit this. I know why this might have happened, but want to make sure it is the expected case.

Test added. It was a bit tricky as the case was only failing in intermittent IR during a pass.

[mravishankar]

Thanks. Thats what I thought. It was for the case where unit-dimensions are all folded away. I wish there was a way to assert that in the code. (I think it is provable that this is always the case based on reshape semantics, but still having an assert that currIndices (in the code above) correspond to all unit-dimensions would be good.

[mravishankar]

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

1141–1142

Actually could you make this !reassociationMap.empty() && !currIndices.empty()

[mravishankar]

Sorry, found a good way to check for the case.

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

1141–1142

Also would be nice to leave a comment like

If the reassociation is empty but the currIndices is not, this by definition is folding unit-dimensions with the result being scalar type. So only append the `currIndices` if reassociation map is not empty.

THat gives me another idea. Could you return failure if the targetShape.size() != 0

[rsuderman]

Fixed for second round of review

[nicolasvasilache]

This one slipped by me a long time ago .. I am concerned about the semantics implications of this change.
For valid collapse / expand ops, the invariant has been that the op needs to list all dimensions in order, from 0 to the largest rank.
It seems the verifier for this is not strict enough.
I see the code is quite old and still based on AffineMap.

I'd like to clean all that old tech debt up and just indices.
Could you please elaborate on what kind of transform needed to introduce such IR ?

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

if (currTargetShape == ShapedType::kDynamicSize &&
       return llvm::None;
 
     // For static shapes, if the product of dimensions of the expanded shape
     // should match the collapsed dimension shape.
     if (prodOfCollapsedDims * sourceShape[sourceDim] != currTargetShape)
       return llvm::None;
 
     currIndices.push_back(sourceDim++);
-    // If there are no dimensions in the target to match, then append the
-    // `currIndices` to the last element of the reassociationMap.
+    // If the reassociation is empty but the currIndices is not, this by
+    // definition is folding unit-dimensions with the result being scalar type.
+    // So only append the `currIndices` if reassociation map is not empty.
     if (targetDim == targetShape.size()) {
+      if (!reassociationMap.empty() && !currIndices.empty())
         reassociationMap.back().append(currIndices.begin(), currIndices.end());
       // Break out of the loops. We should be done here.
       break;
     }
     reassociationMap.emplace_back(ReassociationIndices{});
     std::swap(reassociationMap.back(), currIndices);
     prodOfCollapsedDims = 1;
   }
   // All the dimensions in the two shapes must have been processed.

mlir/test/Dialect/Linalg/canonicalize.mlir

linalg.tiled_loop (%arg3) = (%c0) to (%c192) step (%c24)
     linalg.fill(%16, %cst) : memref<?xf32, #map>, f32
     linalg.yield
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: zero_rank_reshape_multi
+func @zero_rank_reshape_multi(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: return %arg0
+  %0 = linalg.tensor_reshape %arg0 [] : tensor<f32> into tensor<1xf32>
+  %1 = linalg.tensor_reshape %0 [[0, 1]] : tensor<1xf32> into tensor<1x1xf32>
+  %2 = linalg.tensor_reshape %1 [] : tensor<1x1xf32> into tensor<f32>
+  return %2 : tensor<f32>
+}
+
+// -----
+
 func @collapsing_tensor_reshapes(%arg0 : tensor<?x?x?x?x?xf32>) -> tensor<?x?xf32>
 {
   %0 = linalg.tensor_reshape %arg0 [[0, 1], [2], [3, 4]]
       : tensor<?x?x?x?x?xf32> into tensor<?x?x?xf32>
   %1 = linalg.tensor_reshape %0 [[0, 1], [2]]
       : tensor<?x?x?xf32> into tensor<?x?xf32>
   return %1 : tensor<?x?xf32>
 }