[mlir][Linalg] Fix tensor.extract_slice(linalg.init_tensor) canonicalization for rank-reducing extract

[gysit]

LGTM!

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

 struct FoldInitTensorWithExtractSliceOp
     : public OpRewritePattern<tensor::ExtractSliceOp> {
   using OpRewritePattern<tensor::ExtractSliceOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::ExtractSliceOp sliceOp,
                                 PatternRewriter &rewriter) const override {
     if (!sliceOp.source().getDefiningOp<linalg::InitTensorOp>())
       return failure();
+    // ExtractSliceOp may be rank-reducing; its dynamic sizes must be preserved
+    // as well as its result type.
     rewriter.replaceOpWithNewOp<linalg::InitTensorOp>(
         sliceOp, sliceOp.sizes(),
-        llvm::to_vector<4>(llvm::map_range(
-            sliceOp.static_sizes(),
-            [](Attribute attr) { return attr.cast<IntegerAttr>().getInt(); })),
+        sliceOp.result().getType().cast<RankedTensorType>().getShape(),
         sliceOp.getSourceType().getElementType());
     return success();
   }
 };
 
 template <typename TensorReshapeOp>
 struct FoldInitTensorWithTensorReshapeOp
     : public OpRewritePattern<TensorReshapeOp> {

mlir/test/Dialect/Linalg/canonicalize.mlir

func @init_canonicalize(%i : index) {
   %2 = tensor.dim %0, %c1: tensor<?x42xf32>
 
   // CHECK: %[[c42:.*]] = constant 42 : index
   // CHECK: call @some_use(%[[I]], %[[c42]])
   call @some_use(%1, %2) : (index, index) -> ()
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @rank_reducing_init_extract
+func @rank_reducing_init_extract(%sz : index, %idx : index) -> tensor<2xf32> {
+  // CHECK: linalg.init_tensor [2] : tensor<2xf32>
+  %a = linalg.init_tensor [%sz, 2] : tensor<?x2xf32>
+
+  // CHECK-NOT: extract
+  %r = tensor.extract_slice %a[%idx, 0] [1, 2] [1, 1] : tensor<?x2xf32> to tensor<2xf32>
+  return %r: tensor<2xf32>
+}