[mlir][tensor] Add canonicalization of `tensor.extract` from cast-like slice

[christopherbate]

Summary

Adds a canonicalization for tensor.extract. When an extract occurs on
a rank-reducing (of unit-dimensions) tensor.extract_slice operation,
the tensor.extract can act on the source of the tensor.extract_slice.

Depends on D158635

[nicolasvasilache]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

861

While this pattern seems fine to me with this check, I think the other comments in https://reviews.llvm.org/D158643 could still apply?

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
LogicalResult matchAndRewrite(tensor::ExtractOp extract,
     if (!llvm::isa<RankedTensorType>(tensorCast.getSource().getType()))
       return failure();
     rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
         extract, tensorCast.getSource(), extract.getIndices());
     return success();
   }
 };
 
+/// Canonicalizes an extract from a unit-dim-reducing `tensor.extract_slice` op:
+///
+/// %val = tensor.extract_slice %source : : tensor<1x5xi32> to tensor<5xi32>
+/// %extracted_element = tensor.extract %val[%c2] : tensor<i32>
+///
+/// to
+///
+/// %extracted_element = tensor.extract %source[%c0, %c2] : tensor<1x5xi32>
+struct ExtractFromCastLikeSlice : public OpRewritePattern<tensor::ExtractOp> {
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp extract,
+                                PatternRewriter &rewriter) const final {
+    auto sliceOp = extract.getTensor().getDefiningOp<tensor::ExtractSliceOp>();
+    if (!sliceOp)
+      return failure();
+
+    RankedTensorType sourceType = sliceOp.getSourceType();
+    if (!isCastLikeExtractSliceOp(sliceOp))
+      return failure();
+
+    llvm::SmallBitVector droppedDims = sliceOp.getDroppedDims();
+    SmallVector<Value> indices;
+    Value zero = rewriter.create<arith::ConstantIndexOp>(extract.getLoc(), 0);
+    unsigned extractIdx = 0;
+    for (int64_t i = 0, e = sourceType.getRank(); i < e; i++)
+      indices.push_back(droppedDims[i] ? zero
+                                       : extract.getIndices()[extractIdx++]);
+    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(extract, sliceOp.getSource(),
+                                                   indices);
+    return success();
+  }
+};
+
 } // namespace
 
 void ExtractOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "extracted");
 }
 
 LogicalResult ExtractOp::verify() {
if (elementsAttr && elementsAttr.isValidIndex(indices))
       return elementsAttr.getValues<Attribute>()[indices];
   }
 
   return {};
 }
 
 void ExtractOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                             MLIRContext *context) {
-  results.add<ExtractFromTensorCast>(context);
+  results.add<ExtractFromTensorCast, ExtractFromCastLikeSlice>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // FromElementsOp
 //===----------------------------------------------------------------------===//
 
 void FromElementsOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {

mlir/test/Dialect/Tensor/canonicalize.mlir

func.func @extract_from_tensor.cast(%tensor: tensor<9xf32>) -> f32 {
   %c0 = arith.constant 0 : index
   // CHECK-NOT: tensor.cast
   %casted = tensor.cast %tensor : tensor<9xf32> to tensor<?xf32>
   // CHECK-NEXT: tensor.extract %[[TENSOR]][%[[C0]]]
   %result = tensor.extract %casted[%c0] : tensor<?xf32>
   return %result : f32
 }
 
+
+// -----
+
+// CHECK-LABEL: func @extract_from_cast_like_slice
+//  CHECK-SAME: %[[ARG:.+]]: tensor<1x5xf32>
+func.func @extract_from_cast_like_slice(%tensor: tensor<1x5xf32>) -> f32 {
+  %c2 = arith.constant 2 : index
+  %0 = tensor.extract_slice %tensor[0, 0][1, 5][1, 1] : tensor<1x5xf32> to tensor<5xf32>
+  // CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index
+  //     CHECK: %[[EXTRACTED:.+]] = tensor.extract %[[ARG]][%[[C0]], %[[C2]]]
+  %1 = tensor.extract %0[%c2] : tensor<5xf32>
+  // CHECK: return %[[EXTRACTED]]
+  return %1 : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @dont_fold_extract_from_non_cast_like_slice
+//  CHECK-SAME: %[[ARG:.+]]: tensor<1x5xf32>
+func.func @dont_fold_extract_from_non_cast_like_slice(%tensor: tensor<1x5xf32>) -> f32 {
+  %c2 = arith.constant 2 : index
+  // CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[ARG]]
+  %0 = tensor.extract_slice %tensor[0, 0][1, 3][1, 1] : tensor<1x5xf32> to tensor<3xf32>
+  // CHECK: %[[EXTRACTED:.+]] = tensor.extract %[[SLICE]]
+  %1 = tensor.extract %0[%c2] : tensor<3xf32>
+  // CHECK: return %[[EXTRACTED]]
+  return %1 : f32
+}
+
 // -----
 
 // CHECK-LABEL: func @extract_from_tensor.from_elements
 func.func @extract_from_tensor.from_elements(%element : index) -> index {
   // CHECK-SAME: ([[ARG:%.*]]: index)
   %c0 = arith.constant 0 : index
   %tensor = tensor.from_elements %element : tensor<1xindex>
   %extracted_element = tensor.extract %tensor[%c0] : tensor<1xindex>

 <128x128xf32>
   return %unpacked : tensor<128x128xf32>
 }
 
 // -----
 
 // Chain: NC -> NCnc -> NCnc -> NC
 // CHECK: func.func @unpack_pack(
-// CHECK-SAME: %[[T:.+]]: tensor<128x128xf32>, 
+// CHECK-SAME: %[[T:.+]]: tensor<128x128xf32>,
 // CHECK: return %[[T]] : tensor<128x128xf32>
 func.func @unpack_pack(%t: tensor<128x128xf32>, %tile1: index, %tile2: index) -> tensor<128x128xf32> {
   %tensor_empty = tensor.empty(%tile1, %tile2) : tensor<16x16x?x?xf32>
   %packed = tensor.pack %t inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty : tensor<128x128xf32> -> tensor<16x16x?x?xf32>
   %tensor_empty1 = tensor.empty() : tensor<128x128xf32>
   %unpacked = tensor.unpack %packed inner_dims_pos = [0, 1] inner_tiles = [%tile1, %tile2] into %tensor_empty1 : tensor<16x16x?x?xf32> -> tensor
 <128x128xf32>
   return %unpacked : tensor<128x128xf32>