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

[herhut]

Summary

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.

[herhut]

Note that this is already tested by existing tests that use the math dialect.

[pifon2a]

Thank you, Stephan!

[herhut]

also support linalg ops with hidden regions like dot.

[herhut]

clang tidy fix

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

InterfaceMethod<
         OperationState state(
           loc, ConcreteOp::getOperationName(), operands, resultTypes,
           $_op->getAttrs());
         for (Region &r : $_op->getRegions())
           r.cloneInto(state.addRegion(), bvm);
         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.
         Returns a null function if this named op does not define a region
         builder.
       }],
       /*retTy=*/"std::function<void(ImplicitLocOpBuilder &, Block &)>",
       /*methodName=*/"getRegionBuilder",

mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp

 }
 
 /// Create linalg op on buffers given the original tensor-based operation and
 /// the buffers for the outputs.
 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;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization patterns.
 //===----------------------------------------------------------------------===//
 
 namespace {
 

 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     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>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return typeConverter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);

mlir/test/Dialect/Linalg/bufferize.mlir

%read = vector.transfer_read %in[%c0], %cst {in_bounds = [true]}
       : tensor<4xf32>, vector<4xf32>
   %tanh = math.tanh %read : vector<4xf32>
   %write = vector.transfer_write %tanh, %out[%c0] {in_bounds = [true]}
       : vector<4xf32>, tensor<4xf32>
   return
   // 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]]
+}