[mlir] Add canonicalization for the `subtensor` op

[pifon2a]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

3185

wow, that's a long one

mlir/test/Transforms/canonicalize.mlir

1118

super-nit: wrap the line

[nicolasvasilache]

Wrap line

[nicolasvasilache]

Build fix

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = extraBaseClassDeclaration # [{
     /// static representation of offsets, sizes and strides. Special sentinels
     /// encode the dynamic case.
     static Type inferResultType(RankedTensorType sourceRankedTensorType,
                                 ArrayRef<int64_t> staticOffsets,
                                 ArrayRef<int64_t> staticSizes,
                                 ArrayRef<int64_t> staticStrides);
   }];
 
-  // let hasCanonicalizer = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TanhOp
 //===----------------------------------------------------------------------===//
 
 def TanhOp : FloatUnaryOp<"tanh"> {
   let summary = "hyperbolic tangent of the specified value";

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

for (unsigned cstIdx = 0, valIdx = 0, e = constantValues.size();
         continue;
       }
       // Remains dynamic move to next value.
       ++valIdx;
     }
   }
 }
 
+static void replaceWithNewOp(PatternRewriter &rewriter, SubViewOp op,
+                             SubViewOp newOp) {
+  rewriter.replaceOpWithNewOp<MemRefCastOp>(op, newOp, op.getType());
+}
+
+static void replaceWithNewOp(PatternRewriter &rewriter, SubTensorOp op,
+                             SubTensorOp newOp) {
+  rewriter.replaceOpWithNewOp<TensorCastOp>(op, newOp, op.getType());
+}
+
 /// Pattern to rewrite a subview op with constant arguments.
-class SubViewOpConstantArgumentFolder final
-    : public OpRewritePattern<SubViewOp> {
+template <typename OpType>
+class OpWithOffsetSizesAndStridesConstantArgumentFolder final
+    : public OpRewritePattern<OpType> {
 public:
-  using OpRewritePattern<SubViewOp>::OpRewritePattern;
+  using OpRewritePattern<OpType>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(SubViewOp subViewOp,
+  LogicalResult matchAndRewrite(OpType op,
                                 PatternRewriter &rewriter) const override {
     // No constant operand, just return;
-    if (llvm::none_of(subViewOp.getOperands(), [](Value operand) {
+    if (llvm::none_of(op.getOperands(), [](Value operand) {
           return matchPattern(operand, m_ConstantIndex());
         }))
       return failure();
 
     // At least one of offsets/sizes/strides is a new constant.
     // Form the new list of operands and constant attributes from the existing.
-    SmallVector<Value, 8> newOffsets(subViewOp.offsets());
+    SmallVector<Value, 8> newOffsets(op.offsets());
     SmallVector<int64_t, 8> newStaticOffsets =
-        extractFromI64ArrayAttr(subViewOp.static_offsets());
-    assert(newStaticOffsets.size() == subViewOp.getRank());
+        extractFromI64ArrayAttr(op.static_offsets());
+    assert(newStaticOffsets.size() == op.getRank());
     canonicalizeSubViewPart(newOffsets, newStaticOffsets,
                             ShapedType::isDynamicStrideOrOffset);
 
-    SmallVector<Value, 8> newSizes(subViewOp.sizes());
+    SmallVector<Value, 8> newSizes(op.sizes());
     SmallVector<int64_t, 8> newStaticSizes =
-        extractFromI64ArrayAttr(subViewOp.static_sizes());
-    assert(newStaticOffsets.size() == subViewOp.getRank());
+        extractFromI64ArrayAttr(op.static_sizes());
+    assert(newStaticOffsets.size() == op.getRank());
     canonicalizeSubViewPart(newSizes, newStaticSizes, ShapedType::isDynamic);
 
-    SmallVector<Value, 8> newStrides(subViewOp.strides());
+    SmallVector<Value, 8> newStrides(op.strides());
     SmallVector<int64_t, 8> newStaticStrides =
-        extractFromI64ArrayAttr(subViewOp.static_strides());
-    assert(newStaticOffsets.size() == subViewOp.getRank());
+        extractFromI64ArrayAttr(op.static_strides());
+    assert(newStaticOffsets.size() == op.getRank());
     canonicalizeSubViewPart(newStrides, newStaticStrides,
                             ShapedType::isDynamicStrideOrOffset);
 
     // Create the new op in canonical form.
-    auto newSubViewOp = rewriter.create<SubViewOp>(
-        subViewOp.getLoc(), subViewOp.source(), newStaticOffsets,
-        newStaticSizes, newStaticStrides, newOffsets, newSizes, newStrides);
-
-    // Insert a memref_cast for compatibility of the uses of the op.
-    rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, newSubViewOp,
-                                              subViewOp.getType());
+    auto newOp = rewriter.create<OpType>(
+        op.getLoc(), op.source(), newStaticOffsets, newStaticSizes,
+        newStaticStrides, newOffsets, newSizes, newStrides);
+
+    replaceWithNewOp(rewriter, op, newOp);
 
     return success();
   }
 };
 
 } // end anonymous namespace
 
 /// Determines whether MemRefCastOp casts to a more dynamic version of the
rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, subViewOp.getType(),
                                               newSubView);
     return success();
   }
 };
 } // namespace
 
 void SubViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                             MLIRContext *context) {
-  results.insert<SubViewOpConstantArgumentFolder, SubViewOpMemRefCastFolder>(
-      context);
+  results.insert<OpWithOffsetSizesAndStridesConstantArgumentFolder<SubViewOp>,
+                 SubViewOpMemRefCastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, SubTensorOp op) {
   return printOpWithOffsetsSizesAndStrides<SubTensorOp>(p, op);
auto expectedType = SubTensorOp::inferResultType(
       extractFromI64ArrayAttr(op.static_sizes()),
       extractFromI64ArrayAttr(op.static_strides()));
   if (op.getType() != expectedType)
     return op.emitError("expected result type to be ") << expectedType;
 
   return success();
 }
 
+void SubTensorOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                              MLIRContext *context) {
+  results
+      .insert<OpWithOffsetSizesAndStridesConstantArgumentFolder<SubTensorOp>>(
+          context);
+}
+
 //===----------------------------------------------------------------------===//
 // TensorCastOp
 //===----------------------------------------------------------------------===//
 
 bool TensorCastOp::areCastCompatible(Type a, Type b) {
   auto aT = a.dyn_cast<TensorType>();
   auto bT = b.dyn_cast<TensorType>();
   if (!aT || !bT)

mlir/test/Transforms/canonicalize.mlir

 func @tensor_cast_chain_invalid(%input: tensor<4x8xi32>) -> tensor<8x4xi32> {
   // CHECK-NEXT: %[[C1:.*]] = tensor_cast %[[IN]]
   %0 = tensor_cast %input : tensor<4x8xi32> to tensor<?x?xi32>
   // CHECK-NEXT: %[[C2:.*]] = tensor_cast %[[C1]]
   %1 = tensor_cast %0 : tensor<?x?xi32> to tensor<8x4xi32>
   // CHECK-NEXT: return %[[C2]]
   return %1 : tensor<8x4xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @subtensor
+// CHECK-SAME: %[[ARG0:[0-9a-z]*]]: index, %[[ARG1:[0-9a-z]*]]: index
+func @subtensor(%t: tensor<8x16x4xf32>, %arg0 : index, %arg1 : index) -> tensor<?x?x?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %c7 = constant 7 : index
+  %c11 = constant 11 : index
+
+  // CHECK: subtensor %{{.*}}[0, 0, 0] [7, 11, 2] [1, 1, 1] :
+  // CHECK-SAME: tensor<8x16x4xf32> to tensor<7x11x2xf32>
+  // CHECK: tensor_cast %{{.*}} : tensor<7x11x2xf32> to tensor<?x?x?xf32>
+  %1 = subtensor %t[%c0, %c0, %c0] [%c7, %c11, %c2] [%c1, %c1, %c1]
+    : tensor<8x16x4xf32> to tensor<?x?x?xf32>
+
+  // Test: subtensor with one dynamic operand can also be folded.
+  // CHECK: subtensor %{{.*}}[0, 0, 0] [2, %[[ARG0]], 2] [1, 1, 1] :
+  // CHECK-SAME: tensor<?x?x?xf32> to tensor<2x?x2xf32>
+  // CHECK: tensor_cast %{{.*}} : tensor<2x?x2xf32> to tensor<?x?x?xf32>
+  %2 = subtensor %1[%c0, %c0, %c0] [%c2, %arg0, %c2] [%c1, %c1, %c1]
+    : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+
+  return %2 : tensor<?x?x?xf32>
+}