[mlir][tensor] Add pattern to rewrite tensor.generate as a constant

[springerm]

Summary

Only ops with a static tensor type and a constant yield value are rewritten.

Depends On: D152455

mlir/include/mlir/Dialect/Tensor/TransformOps/TensorTransformOps.td

let description = [{
     Indicates that reassociative reshapes (tensor.collapse_shape /
     tensor.expand_shape) should be folded with inverse rank expansions / rank
     reductions (via tensor.insert_slice / tensor.extract_slice).
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyRewriteTensorOpsAsConstantPatternsOp : Op<Transform_Dialect,
+    "apply_patterns.tensor.rewrite_as_constant",
+    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+  let description = [{
+    Indicates that tensor ops (such as tensor.generate) should be replaced with
+    constants (arith.constant) when possible.
+  }];
+
+  let assemblyFormat = "attr-dict";
+}
+
 def Transform_TensorPadOp : Transform_ConcreteOpType<"tensor.pad">;
 
 def MakeLoopIndependentOp
     : Op<Transform_Dialect, "tensor.make_loop_independent",
          [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
           TransformOpInterface, TransformEachOpTrait]> {
   let description = [{
     Rewrite the targeted ops such that their index-typed operands no longer

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

 void populateFoldTensorEmptyPatterns(RewritePatternSet &patterns,
                                      bool foldSingleUseOnly = false);
 
 /// Populates `patterns` with patterns that fold operations like `tensor.pad`
 /// and `tensor.extract_slice` into `tensor.pack` and `tensor.unpack` operations
 /// respectively.
 void populateFoldIntoPackAndUnpackPatterns(RewritePatternSet &patterns);
 
+/// Populates `patterns` with patterns that replace tensor ops (such as
+/// tensor.generate) with constants when possible.
+void populateRewriteAsConstantPatterns(RewritePatternSet &patterns);
+
 //===----------------------------------------------------------------------===//
 // Transform helpers
 //===----------------------------------------------------------------------===//
 
 /// Build a new tensor::PadOp with low/high padding that is independent of all
 /// given independencies. If the op is already independent of all
 /// independencies, the same PadOp result is returned.
 ///

mlir/lib/Dialect/Tensor/TransformOps/TensorTransformOps.cpp

void transform::ApplyMergeConsecutiveInsertExtractSlicePatternsOp::
   tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
 }
 
 void transform::ApplyReassociativeReshapeFoldingPatternsOp::populatePatterns(
     RewritePatternSet &patterns) {
   tensor::populateReassociativeReshapeFoldingPatterns(patterns);
 }
 
+void transform::ApplyRewriteTensorOpsAsConstantPatternsOp::populatePatterns(
+    RewritePatternSet &patterns) {
+  tensor::populateRewriteAsConstantPatterns(patterns);
+}
+
 //===----------------------------------------------------------------------===//
 // MakeLoopIndependentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::MakeLoopIndependentOp::applyToOne(
     Operation *target, transform::ApplyToEachResultList &results,
     transform::TransformState &state) {
   // Gather IVs.

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
   EmptyOpPatterns.cpp
   ExtractSliceFromReshapeUtils.cpp
   FoldIntoPackAndUnpackPatterns.cpp
   FoldTensorSubsetOps.cpp
   IndependenceTransforms.cpp
   MergeConsecutiveInsertExtractSlicePatterns.cpp
   ReshapePatterns.cpp
+  RewriteAsConstant.cpp
   SwapExtractSliceWithProducerPatterns.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 

mlir/lib/Dialect/Tensor/Transforms/RewriteAsConstant.cpp

+//===- RewriteAsConstant.cpp - Patterns to rewrite tensor ops as constants ===//
+//
+// 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/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/PatternMatch.h"
+
+using namespace mlir;
+using namespace mlir::tensor;
+
+namespace {
+
+/// Rewrite tensor.generate with arith.constant if the yielded value is a
+/// constant and the tensor type is static.
+struct GenerateToConstant : public OpRewritePattern<GenerateOp> {
+  using OpRewritePattern<GenerateOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenerateOp generateOp,
+                                PatternRewriter &rewriter) const override {
+    auto tensorType =
+        llvm::cast<RankedTensorType>(generateOp.getResult().getType());
+    if (!tensorType.hasStaticShape())
+      return failure();
+    auto terminatorOp =
+        cast<tensor::YieldOp>(generateOp.getBody().front().getTerminator());
+    Attribute attr;
+    if (!matchPattern(terminatorOp.getValue(), m_Constant(&attr)))
+      return failure();
+    Operation *constantOp =
+        rewriter.getContext()
+            ->getLoadedDialect<TensorDialect>()
+            ->materializeConstant(rewriter,
+                                  DenseElementsAttr::get(tensorType, attr),
+                                  tensorType, generateOp->getLoc());
+    if (!constantOp)
+      return failure();
+    rewriter.replaceOp(generateOp, constantOp->getResults());
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::tensor::populateRewriteAsConstantPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<GenerateToConstant>(patterns.getContext());
+}

mlir/test/Dialect/Tensor/rewrite-as-constant.mlir

+// RUN: mlir-opt -split-input-file -test-transform-dialect-interpreter %s | FileCheck %s
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !transform.any_op):
+  transform.apply_patterns to %module_op {
+    transform.apply_patterns.tensor.rewrite_as_constant
+  } : !transform.any_op
+}
+
+// CHECK-LABEL: func @tensor_generate_constant(
+//       CHECK:   %[[cst:.*]] = arith.constant dense<5.000000e+00> : tensor<2x3x5xf32>
+//       CHECK:   return %[[cst]]
+func.func @tensor_generate_constant() -> tensor<2x3x5xf32> {
+  %cst = arith.constant 5.0 : f32
+  %0 = tensor.generate {
+    ^bb0(%arg0: index, %arg1: index, %arg2: index):
+    tensor.yield %cst : f32
+  } : tensor<2x3x5xf32>
+  return %0 : tensor<2x3x5xf32>
+}