[mlir][Standard] Add canonicalizer for dynamic_tensor_from_elements

[herhut]

Summary

This adds canonicalizer for

• extracting an element from a dynamic_tensor_from_elements
• propagating constant operands to the type of dynamic_tensor_from_elements

[mehdi_amini]

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

1739

"resulting" does not read well to me here?

1756

can you specify the type here?

1803

auto tensorFromElements = extract.aggregate().getDefiningOp<DynamicTensorFromElementsOp>();

1818

Is this always a more canonical form to do this?

[bondhugula]

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

1736–1739

Triple comment please.

1786–1795

Triple /// comments please.

[frgossen]

Thanks!

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

1818

If the element is extracted multiple times from the same tensor you can end up with duplicate computations here.
I think this one should not be a general canonicalisation pattern.

mlir/test/Transforms/canonicalize.mlir

998

can we have another test for rank > 1?

[herhut]

Comments...

[herhut]

More comments

[herhut]

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

1818

Yes, I think this is always a better form. It might give you some code-growth if you have many uses (as in, many extract element operations) but you avoid an intermediate data structure, which generally propagates knowledge and makes other patterns easier to write.

I could guard this with the number of use sites being small to avoid excessive code growth but I think that would be pathological cases. Also dynamic_tensor_from_elements is not really a "compute op" like a LinAlg generic but is meant to create small intermediate tensors where pattern composition requires this, like in the shape dialect. I would not write the same pattern for LinAlg for instance.

I have added a check that there are no side-effects to make sure that replicating computation is OK. If you insist, I can also add a check that it is a single use. However, then we might get odd knowledge propagation in cases where one first does a CSE to combine for example shape_of computations vs. doing that later.

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

def DynamicTensorFromElementsOp : Std_Op<"dynamic_tensor_from_elements",
   let regions = (region SizedRegion<1>:$body);
 
   let builders = [
     // Build op and populate its body per callback function.
     OpBuilder<"OpBuilder &b, OperationState &result, Type resultTy, "
               "ValueRange dynamicExtents, "
               "function_ref<void(OpBuilder &, Location, ValueRange)>">,
   ];
+
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ExpOp
 //===----------------------------------------------------------------------===//
 
 def ExpOp : FloatUnaryOp<"exp"> {
   let summary = "base-e exponential of the specified value";

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

 //===- Ops.cpp - Standard MLIR Operations ---------------------------------===//
 //
 // 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/StandardOps/IR/Ops.h"
 
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/IR/TypeUtilities.h"
void DynamicTensorFromElementsOp::build(
   Region *bodyRegion = result.regions.front().get();
   auto rank = resultTy.cast<RankedTensorType>().getRank();
   SmallVector<Type, 2> argumentTypes(rank, b.getIndexType());
   Block *bodyBlock =
       b.createBlock(bodyRegion, bodyRegion->end(), argumentTypes);
   bodyBuilder(b, result.location, bodyBlock->getArguments());
 }
 
+namespace {
+
+/// Canonicalizes dynamic_tensor_from_elements operations with a constant
+/// operand into the equivalent operation with the operand expressed in the
+/// result type, instead. We also insert a type cast to make sure that the
+/// resulting IR is still well-typed.
+struct StaticDynamicTensorFromElements
+    : public OpRewritePattern<DynamicTensorFromElementsOp> {
+  using OpRewritePattern<DynamicTensorFromElementsOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DynamicTensorFromElementsOp tensorFromElements,
+                                PatternRewriter &rewriter) const final {
+    auto resultType =
+        tensorFromElements.getResult().getType().cast<RankedTensorType>();
+
+    if (resultType.hasStaticShape())
+      return failure();
+
+    SmallVector<Value, 4> newOperands;
+    SmallVector<int64_t, 4> newShape;
+    auto operandsIt = tensorFromElements.dynamicExtents().begin();
+
+    for (int64_t dim : resultType.getShape()) {
+      if (dim != RankedTensorType::kDynamicSize) {
+        newShape.push_back(dim);
+        continue;
+      }
+      APInt index;
+      if (!matchPattern(*operandsIt, m_ConstantInt(&index))) {
+        newShape.push_back(RankedTensorType::kDynamicSize);
+        newOperands.push_back(*operandsIt++);
+        continue;
+      }
+      newShape.push_back(index.getSExtValue());
+      operandsIt++;
+    }
+
+    if (newOperands.size() == tensorFromElements.dynamicExtents().size())
+      return failure();
+
+    auto loc = tensorFromElements.getLoc();
+    auto newOp = rewriter.create<DynamicTensorFromElementsOp>(
+        loc, RankedTensorType::get(newShape, resultType.getElementType()),
+        newOperands);
+    rewriter.inlineRegionBefore(tensorFromElements.body(), newOp.body(),
+                                newOp.body().begin());
+    rewriter.replaceOpWithNewOp<TensorCastOp>(tensorFromElements, resultType,
+                                              newOp);
+    return success();
+  }
+};
+
+/// Canonicalizes the pattern of the form
+///
+/// %tensor = dynamic_tensor_from_elements %x {
+///   ^bb0(%arg0: index):  // no predecessors
+///   <computation>
+///   yield %1 : index
+/// } : tensor<?xindex>
+/// %extracted_element = extract_element %tensor[%c0] : tensor<?xi32>
+///
+/// to just <computation> with %arg0 replaced by %c0. We only do this if the
+/// dynamic_tensor_from_elements operation has no side-effects.
+struct ExtractElementFromDynamicTensorFromElements
+    : public OpRewritePattern<ExtractElementOp> {
+  using OpRewritePattern<ExtractElementOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ExtractElementOp extract,
+                                PatternRewriter &rewriter) const final {
+    auto tensorFromElements =
+        extract.aggregate().getDefiningOp<DynamicTensorFromElementsOp>();
+    if (!tensorFromElements || !wouldOpBeTriviallyDead(tensorFromElements))
+      return failure();
+
+    BlockAndValueMapping mapping;
+    Block *body = tensorFromElements.getBody();
+    mapping.map(body->getArguments(), extract.indices());
+    for (auto &op : body->without_terminator())
+      rewriter.clone(op, mapping);
+
+    auto yield = cast<YieldOp>(body->getTerminator());
+
+    rewriter.replaceOp(extract, mapping.lookupOrDefault(yield.value()));
+    return success();
+  }
+};
+
+} // namespace
+
+void DynamicTensorFromElementsOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &results, MLIRContext *context) {
+  results.insert<ExtractElementFromDynamicTensorFromElements,
+                 StaticDynamicTensorFromElements>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ExtractElementOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ExtractElementOp op) {
   // Verify the # indices match if we have a ranked type.
   auto aggregateType = op.getAggregate().getType().cast<ShapedType>();
   if (aggregateType.hasRank() &&
struct ExtractElementFromTensorFromElements
     : public OpRewritePattern<ExtractElementOp> {
   using OpRewritePattern<ExtractElementOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ExtractElementOp extract,
                                 PatternRewriter &rewriter) const final {
     if (extract.indices().size() != 1)
       return failure();
 
-    auto tensor_from_elements = dyn_cast_or_null<TensorFromElementsOp>(
+    auto tensorFromElements = dyn_cast_or_null<TensorFromElementsOp>(
         extract.aggregate().getDefiningOp());
-    if (tensor_from_elements == nullptr)
+    if (tensorFromElements == nullptr)
       return failure();
 
     APInt index;
     if (!matchPattern(*extract.indices().begin(), m_ConstantInt(&index)))
       return failure();
     rewriter.replaceOp(extract,
-                       tensor_from_elements.getOperand(index.getZExtValue()));
+                       tensorFromElements.getOperand(index.getZExtValue()));
     return success();
   }
 };
 
 } // namespace
 
 void TensorFromElementsOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {

mlir/test/Transforms/canonicalize.mlir

 func @extract_element_from_tensor_from_elements(%element : index) -> index {
   // CHECK-SAME: ([[ARG:%.*]]: index)
   %c0 = constant 0 : index
   %tensor = tensor_from_elements %element : tensor<1xindex>
   %extracted_element = extract_element %tensor[%c0] : tensor<1xindex>
   // CHECK: [[ARG]] : index
   return %extracted_element : index
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements
+// CHECK-SAME: %[[IDX:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>
+func @extract_element_from_dynamic_tensor_from_elements(%idx: index, %tensor: tensor<*xf32>) -> index {
+  %size = rank %tensor : tensor<*xf32>
+  // CHECK-NEXT: %[[RES:.*]] = dim %[[TENSOR]], %[[IDX]]
+  %0 = dynamic_tensor_from_elements %size {
+    ^bb0(%arg0: index):
+    %1 = dim %tensor, %arg0 : tensor<*xf32>
+    yield %1 : index
+  } : tensor<?xindex>
+  %1 = extract_element %0[%idx] : tensor<?xindex>
+  // CHECK-NEXT: return %[[RES]]
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements_2d
+// CHECK-SAME: %[[IDX0:.*]]: index, %[[IDX1:.*]]: index, %[[TENSOR:.*]]: tensor<*xf32>
+func @extract_element_from_dynamic_tensor_from_elements_2d(%idx0: index, %idx1: index, %tensor: tensor<*xf32>) -> index {
+  %size = rank %tensor : tensor<*xf32>
+  // CHECK-NEXT: %[[DIM0:.*]] = dim %[[TENSOR]], %[[IDX0]]
+  // CHECK-NEXT: %[[DIM1:.*]] = dim %[[TENSOR]], %[[IDX1]]
+  // CHECK-NEXT: %[[RES:.*]] = addi %[[DIM0]], %[[DIM1]]
+  %0 = dynamic_tensor_from_elements %size, %size {
+    ^bb0(%arg0: index, %arg1: index):
+    %1 = dim %tensor, %arg0 : tensor<*xf32>
+    %2 = dim %tensor, %arg1 : tensor<*xf32>
+    %3 = addi %1, %2 : index
+    yield %3 : index
+  } : tensor<?x?xindex>
+  %4 = extract_element %0[%idx0, %idx1] : tensor<?x?xindex>
+  // CHECK-NEXT: return %[[RES]]
+  return %4 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_element_from_dynamic_tensor_from_elements_sideeffects
+// CHECK-SAME: %[[IDX:.*]]: index
+func @extract_element_from_dynamic_tensor_from_elements_sideeffects(%idx: index, %tensor: tensor<*xf32>) -> index {
+  %size = rank %tensor : tensor<*xf32>
+  %mem = alloc(%size) : memref<?xindex>
+  // CHECK: %[[DTENSOR:.*]] = dynamic_tensor_from_elements
+  %0 = dynamic_tensor_from_elements %size {
+    ^bb0(%arg0: index):
+    %1 = dim %tensor, %arg0 : tensor<*xf32>
+    store %1, %mem[%arg0] : memref<?xindex>
+    yield %1 : index
+  } : tensor<?xindex>
+  // CHECK: %[[RES:.*]] = extract_element %[[DTENSOR]][%[[IDX]]]
+  %1 = extract_element %0[%idx] : tensor<?xindex>
+  // CHECK-NEXT: return %[[RES]]
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: @static_dynamic_tensor_from_elements
+// CHECK-SAME: %[[SIZE1:.*]]: index, %[[SIZE4:.*]]: index)
+func @static_dynamic_tensor_from_elements(%size1: index, %size4: index) -> tensor<3x?x?x7x?xindex> {
+  %c5 = constant 5 : index
+  // CHECK: dynamic_tensor_from_elements %[[SIZE1]], %[[SIZE4]]
+  %0 = dynamic_tensor_from_elements %size1, %c5, %size4 {
+    ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index):
+    %1 = constant 32 : index
+    yield %1 : index
+  // CHECK: : tensor<3x?x5x7x?xindex>
+  } : tensor<3x?x?x7x?xindex>
+  // CHECK: tensor_cast %{{.*}} : tensor<3x?x5x7x?xindex> to tensor<3x?x?x7x?xindex>
+  return %0 : tensor<3x?x?x7x?xindex>
+}
+