[mlir] Use static shape knowledge when lowering memref.reshape

[bkramer]

Summary

This is actually necessary for correctness, as memref.reinterpret_cast
doesn't verify if the output shape doesn't match the static sizes.

[pifon2a]

Thanks, Ben!

mlir/lib/Dialect/StandardOps/Transforms/ExpandOps.cpp

LogicalResult matchAndRewrite(memref::ReshapeOp op,
     int64_t rank = shapeType.cast<MemRefType>().getDimSize(0);
     SmallVector<OpFoldResult, 4> sizes, strides;
     sizes.resize(rank);
     strides.resize(rank);
 
     Location loc = op.getLoc();
     Value stride = rewriter.create<ConstantIndexOp>(loc, 1);
     for (int i = rank - 1; i >= 0; --i) {
+      Value size;
+      // Load dynamic sizes from the shape input, use constants for static dims.
+      if (op.getType().isDynamicDim(i)) {
         Value index = rewriter.create<ConstantIndexOp>(loc, i);
-      Value size = rewriter.create<memref::LoadOp>(loc, op.shape(), index);
+        size = rewriter.create<memref::LoadOp>(loc, op.shape(), index);
         if (!size.getType().isa<IndexType>())
-        size = rewriter.create<IndexCastOp>(loc, size, rewriter.getIndexType());
+          size =
+              rewriter.create<IndexCastOp>(loc, size, rewriter.getIndexType());
         sizes[i] = size;
+      } else {
+        sizes[i] = rewriter.getIndexAttr(op.getType().getDimSize(i));
+        size = rewriter.create<ConstantOp>(loc, sizes[i].get<Attribute>());
+      }
       strides[i] = stride;
       if (i > 0)
         stride = rewriter.create<MulIOp>(loc, stride, size);
     }
     rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
         op, op.getType(), op.source(), /*offset=*/rewriter.getIndexAttr(0),
         sizes, strides);
     return success();

mlir/test/Dialect/Standard/expand-ops.mlir

 // CHECK:           [[CMP2:%.+]] = or [[TERM1]], [[TERM2]] : i1
 // CHECK:           [[RES:%.+]] = select [[CMP2]], [[TRUE3]], [[FALSE]] : i32
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_reshape(
 func @memref_reshape(%input: memref<*xf32>,
-                     %shape: memref<3xi32>) -> memref<?x?x?xf32> {
+                     %shape: memref<3xi32>) -> memref<?x?x8xf32> {
   %result = memref.reshape %input(%shape)
-               : (memref<*xf32>, memref<3xi32>) -> memref<?x?x?xf32>
-  return %result : memref<?x?x?xf32>
+               : (memref<*xf32>, memref<3xi32>) -> memref<?x?x8xf32>
+  return %result : memref<?x?x8xf32>
 }
 // CHECK-SAME: [[SRC:%.*]]: memref<*xf32>,
-// CHECK-SAME: [[SHAPE:%.*]]: memref<3xi32>) -> memref<?x?x?xf32> {
+// CHECK-SAME: [[SHAPE:%.*]]: memref<3xi32>) -> memref<?x?x8xf32> {
 
 // CHECK: [[C1:%.*]] = constant 1 : index
-// CHECK: [[C2:%.*]] = constant 2 : index
-// CHECK: [[DIM_2:%.*]] = memref.load [[SHAPE]]{{\[}}[[C2]]] : memref<3xi32>
-// CHECK: [[SIZE_2:%.*]] = index_cast [[DIM_2]] : i32 to index
-// CHECK: [[STRIDE_1:%.*]] = muli [[C1]], [[SIZE_2]] : index
+// CHECK: [[C8:%.*]] = constant 8 : index
+// CHECK: [[STRIDE_1:%.*]] = muli [[C1]], [[C8]] : index
 
 // CHECK: [[C1_:%.*]] = constant 1 : index
 // CHECK: [[DIM_1:%.*]] = memref.load [[SHAPE]]{{\[}}[[C1_]]] : memref<3xi32>
 // CHECK: [[SIZE_1:%.*]] = index_cast [[DIM_1]] : i32 to index
 // CHECK: [[STRIDE_0:%.*]] = muli [[STRIDE_1]], [[SIZE_1]] : index
 
 // CHECK: [[C0:%.*]] = constant 0 : index
 // CHECK: [[DIM_0:%.*]] = memref.load [[SHAPE]]{{\[}}[[C0]]] : memref<3xi32>
 // CHECK: [[SIZE_0:%.*]] = index_cast [[DIM_0]] : i32 to index
 
 // CHECK: [[RESULT:%.*]] = memref.reinterpret_cast [[SRC]]
-// CHECK-SAME: to offset: [0], sizes: {{\[}}[[SIZE_0]], [[SIZE_1]], [[SIZE_2]]],
+// CHECK-SAME: to offset: [0], sizes: {{\[}}[[SIZE_0]], [[SIZE_1]], 8],
 // CHECK-SAME: strides: {{\[}}[[STRIDE_0]], [[STRIDE_1]], [[C1]]]
-// CHECK-SAME: : memref<*xf32> to memref<?x?x?xf32>
+// CHECK-SAME: : memref<*xf32> to memref<?x?x8xf32>