[MLIR][Bufferization] Introduce `EmptyTensorToAllocTensorOp`

[chelini]

Summary

Introduce a new transform operation to replace tensor.empty with
alloc_tensor operations. The operation is a pass-through if the target
operation is already a alloc_tensor; otherwise, it expects a
tensor.empty as a target. Currently, it does not return any results.

The operation is expected to run before one_shot_bufferize as
one_shot_bufferize rejects tensor.empty.

[springerm]

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td

75–77

• Consumes the target handle, produces the result handle
• target must be mapped to only tensor.empty ops
• always succeeds.

81

You can try using Transform_ConcreteOpType<"tensor.empty"> here. See LoopPeelOp in SCFTransformOps.td for an example. It's a typed alternative for PDL_Operation. We don't use it widely yet, but here it would fit well.

82

Add result for the generated alloc_tensor op here.

88

You can write ::mlir::tensor::EmptyOp here, then you also don't need check the type of the op in the impl.

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

81

IRRewriter

[chelini]

• Use Transform_ConcreteOpType for target and transformed handle.

• Add result

• Use IRRewriter

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h

 //===- BufferizationTransformOps.h - Buff. transf. ops ----------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
+#include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
+namespace mlir {
+namespace tensor {
+class EmptyOp;
+} // namespace tensor
+} // namespace mlir
+
 //===----------------------------------------------------------------------===//
 // Bufferization Transform Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.h.inc"
 
 namespace mlir {
 class DialectRegistry;
 
 namespace bufferization {
 void registerTransformDialectExtension(DialectRegistry &registry);
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMOPS_BUFFERIZATIONTRANSFORMOPS_H

mlir/include/mlir/Dialect/Bufferization/TransformOps/BufferizationTransformOps.td

 //===- BufferizationTransformOps.td - Buff. transf. ops ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BUFFERIZATION_TRANSFORM_OPS
 #define BUFFERIZATION_TRANSFORM_OPS
 
 include "mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
 include "mlir/Dialect/Transform/IR/TransformDialect.td"
 include "mlir/Dialect/Transform/IR/TransformEffects.td"
 include "mlir/Dialect/Transform/IR/TransformInterfaces.td"
+include "mlir/Dialect/Transform/IR/TransformTypes.td"
 include "mlir/Dialect/PDL/IR/PDLTypes.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/OpBase.td"
 
 def OneShotBufferizeOp
     : Op<Transform_Dialect, "bufferization.one_shot_bufferize",
         [DeclareOpInterfaceMethods<TransformOpInterface>,
          DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
def OneShotBufferizeOp
   let results = (outs);
 
   let assemblyFormat = [{
     (`layout` `{` $function_boundary_type_conversion^ `}`)?
     $target attr-dict
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// EmptyTensorToAllocTensorOp
+//===----------------------------------------------------------------------===//
+
+
+def Transform_EmptyOp : Transform_ConcreteOpType<"tensor.empty">;
+def Transform_AllocTensorOp : Transform_ConcreteOpType<"bufferization.alloc_tensor">;
+
+def EmptyTensorToAllocTensorOp 
+    : Op<Transform_Dialect, "bufferization.empty_tensor_to_alloc_tensor",
+        [FunctionalStyleTransformOpTrait, 
+         MemoryEffectsOpInterface,
+         TransformOpInterface,
+         TransformEachOpTrait]> {
+  let description = [{
+    Replace a tensor.empty with a bufferization.tensor_alloc.
+    
+    ### Return modes
+
+    This operation consumes the `target` handle and produces the `transformed`
+    handle. `target` is expected to be a `tensor.empty` operation. The transform
+    always succeeds.
+  }];
+
+  let arguments = (ins Transform_EmptyOp:$target);
+  let results = (outs Transform_AllocTensorOp:$transformed);
+
+  let assemblyFormat = "$target attr-dict `:` functional-type(operands, results)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::tensor::EmptyOp target,
+        ::llvm::SmallVector<::mlir::Operation *> &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // BUFFERIZATION_TRANSFORM_OPS

mlir/lib/Dialect/Bufferization/TransformOps/BufferizationTransformOps.cpp

 //===- BufferizationTransformOps.h - Bufferization transform ops ----------===//
 //
 // 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/Bufferization/TransformOps/BufferizationTransformOps.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Transform/IR/TransformDialect.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::transform;
 
 //===----------------------------------------------------------------------===//
 // OneShotBufferizeOp
void transform::OneShotBufferizeOp::getEffects(
   effects.emplace_back(MemoryEffects::Read::get(), getTarget(),
                        TransformMappingResource::get());
 
   // Handles that are not modules are not longer usable.
   if (!getTargetIsModule())
     effects.emplace_back(MemoryEffects::Free::get(), getTarget(),
                          TransformMappingResource::get());
 }
+
+//===----------------------------------------------------------------------===//
+// EmptyTensorToAllocTensorOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+EmptyTensorToAllocTensorOp::applyToOne(tensor::EmptyOp target,
+                                       SmallVector<Operation *> &results,
+                                       transform::TransformState &state) {
+  IRRewriter rewriter(target->getContext());
+  rewriter.setInsertionPoint(target);
+  auto alloc = rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
+      target, target.getType(), target.getDynamicSizes());
+  results.push_back(alloc);
+  return DiagnosedSilenceableFailure::success();
+}
+
 //===----------------------------------------------------------------------===//
 // Transform op registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// Registers new ops and declares PDL as dependent dialect since the additional
 /// ops are using PDL types for operands and results.
 class BufferizationTransformDialectExtension

mlir/test/Dialect/Bufferization/Transforms/transform-ops.mlir

 // CHECK-SAME:  %[[B:.*]]: memref<9x6xf32>,
 // CHECK-SAME:  %[[C:.*]]: memref<12x6xf32>) -> memref<12x6xf32> {
 func.func @matmul(%A: tensor<12x9xf32>, %B: tensor<9x6xf32>, %C: tensor<12x6xf32>) -> tensor<12x6xf32> {
   // CHECK: linalg.matmul ins(%[[A]], %[[B]] : memref<12x9xf32>, memref<9x6xf32>) outs(%[[C]] : memref<12x6xf32>)
   %D = linalg.matmul ins(%A, %B: tensor<12x9xf32>, tensor<9x6xf32>) outs(%C: tensor<12x6xf32>) -> tensor<12x6xf32>
   // CHECK: return %[[C]] : memref<12x6xf32>
   return %D : tensor<12x6xf32>
 }
+
+// -----
+
+transform.sequence failures(propagate) {
+  ^bb0(%arg1: !pdl.operation):
+    %0 = transform.structured.match ops{["tensor.empty"]} in %arg1
+    %1 = transform.cast %0 : !pdl.operation to !transform.op<"tensor.empty">
+    transform.bufferization.empty_tensor_to_alloc_tensor %1 : (!transform.op<"tensor.empty">) -> !transform.op<"bufferization.alloc_tensor">
+}
+
+// Expect `bufferization.empty_tensor_to_alloc_tensor` to replace the tensor.empty.
+func.func @empty_to_tensor_alloc() -> tensor<2x2xf32> {
+  // CHECK: bufferization.alloc_tensor
+  %0 = tensor.empty() : tensor<2x2xf32>
+  return %0 : tensor<2x2xf32>
+}