[Mlir-commits] [mlir] 5e6c170 - [mlir][linalg] Fix bufferize pattern to allow unknown operations in body of generic

[Stephan Herhut]

Author: Stephan Herhut
Date: 2021-09-10T13:37:42+02:00
New Revision: 5e6c170b3f41f4da3888a0c4e09eb6af06d7d2c9

URL: https://github.com/llvm/llvm-project/commit/5e6c170b3f41f4da3888a0c4e09eb6af06d7d2c9
DIFF: https://github.com/llvm/llvm-project/commit/5e6c170b3f41f4da3888a0c4e09eb6af06d7d2c9.diff

LOG: [mlir][linalg] Fix bufferize pattern to allow unknown operations in body of generic

The original version of the bufferization pattern for linalg.generic would
manually clone operations within the region to the bufferized clone of the
operation. This triggers legality requirements on those operations in the
conversion infra. Instead, this now uses the rewriter to inline the region
instead, avoiding those legality requirements.

Differential Revision: https://reviews.llvm.org/D109581

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
    mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
    mlir/test/Dialect/Linalg/bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 36146c341f68c..fe8ddaffbc169 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -890,6 +890,27 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
         return b.createOperation(state);
       }]
     >,
+    InterfaceMethod<
+      /*desc=*/[{
+        Clone the current operation with the given location, operands
+        and BlockAndValueMapping but leave the regions empty. This is
+        used to abstract away the optional underlying region creation.
+        This does not change the balance between input, output_buffer
+        and init_tensors operands.
+      }],
+      /*retTy=*/"Operation *",
+      /*methodName=*/"cloneWithoutRegions",
+      (ins "OpBuilder &":$b, "Location":$loc, "TypeRange":$resultTypes,
+           "ValueRange":$operands),
+      [{
+        OperationState state(
+          loc, ConcreteOp::getOperationName(), operands, resultTypes,
+          $_op->getAttrs());
+        for (size_t cnt = 0, e = $_op->getNumRegions(); cnt < e; ++cnt)
+          state.addRegion();
+        return b.createOperation(state);
+      }]
+    >,
     StaticInterfaceMethod<
       /*desc=*/[{
         Returns the region builder for constructing the body for linalg.generic.

diff  --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
index c660b32eca35f..60ccd0bd787d3 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp
@@ -79,38 +79,17 @@ LinalgOp
 mlir::linalg::createLinalgOpOnBuffers(ConversionPatternRewriter &rewriter,
                                       LinalgOp linalgOp, ValueRange inputs,
                                       ValueRange outputs) {
-  if (auto genericOp = mlir::dyn_cast<GenericOp>(*linalgOp)) {
-    // Generate a new linalg operation that works on buffers.
-    auto newGenericOp = rewriter.create<GenericOp>(
-        genericOp.getLoc(),
-        /*resultTensorTypes=*/llvm::None,
-        /*inputs=*/inputs,
-        /*outputs=*/outputs, genericOp.indexing_maps(),
-        genericOp.iterator_types(), genericOp.docAttr(),
-        genericOp.library_callAttr());
-
-    // Create a new block in the region of the new Generic Op.
-    Block *oldBlock = genericOp.getBody();
-    Region &newRegion = newGenericOp.region();
-    Block *newBlock = rewriter.createBlock(&newRegion, newRegion.begin(),
-                                           oldBlock->getArgumentTypes());
-
-    // Clone the body of the old block to the new block.
-    BlockAndValueMapping mapping;
-    mapping.map(oldBlock->getArguments(), newBlock->getArguments());
-
-    OpBuilder::InsertionGuard guard(rewriter);
-    rewriter.setInsertionPointToEnd(newBlock);
-    for (auto &op : oldBlock->getOperations()) {
-      Operation *clonedOp = rewriter.clone(op, mapping);
-      mapping.map(op.getResults(), clonedOp->getResults());
-    }
-    return newGenericOp;
-  }
   SmallVector<Value, 8> newOperands = inputs;
   newOperands.append(outputs.begin(), outputs.end());
-  return linalgOp.clone(rewriter, linalgOp.getLoc(),
-                        /*resultTypes=*/ArrayRef<Type>{}, newOperands);
+  auto *newOp = linalgOp.cloneWithoutRegions(rewriter, linalgOp.getLoc(),
+                                             /*resultTypes=*/ArrayRef<Type>{},
+                                             newOperands);
+  for (auto regions : llvm::zip(linalgOp->getRegions(), newOp->getRegions())) {
+    auto &oldRegion = std::get<0>(regions);
+    auto &newRegion = std::get<1>(regions);
+    rewriter.inlineRegionBefore(oldRegion, newRegion, newRegion.begin());
+  }
+  return newOp;
 }
 
 //===----------------------------------------------------------------------===//
@@ -344,9 +323,8 @@ struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
     BufferizeTypeConverter typeConverter;
 
     // Mark all Standard operations legal.
-    target.addLegalDialect<AffineDialect, math::MathDialect,
-                           memref::MemRefDialect, StandardOpsDialect,
-                           tensor::TensorDialect>();
+    target.addLegalDialect<AffineDialect, memref::MemRefDialect,
+                           StandardOpsDialect, tensor::TensorDialect>();
     target.addIllegalOp<InitTensorOp, tensor::ExtractSliceOp,
                         tensor::InsertSliceOp, PadTensorOp>();
 

diff  --git a/mlir/test/Dialect/Linalg/bufferize.mlir b/mlir/test/Dialect/Linalg/bufferize.mlir
index 86a3f2cbc5cc8..1bf709824a034 100644
--- a/mlir/test/Dialect/Linalg/bufferize.mlir
+++ b/mlir/test/Dialect/Linalg/bufferize.mlir
@@ -316,3 +316,16 @@ func @vector_transfer(%in: tensor<4xf32>, %out: tensor<4xf32>) {
   // CHECK: vector.transfer_read {{.*}} : memref<4xf32>, vector<4xf32>
   // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, memref<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL:   func @bufferize_dot
+func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> {
+  %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>)
+                          outs(%out : tensor<f32>) -> tensor<f32>
+  return %dot : tensor<f32>
+  // CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>)
+  // CHECK-SAME:       outs(%[[OUT:.*]] : memref<f32>)
+  // CHECK: %[[OUT_TENSOR:.*]] = memref.tensor_load %[[OUT]] : memref<f32>
+  // CHECK: return %[[OUT_TENSOR]]
+}