[mlir] Remove clone methods from DPS interface.

[pifon2a]

Typo

[pifon2a]

Clean-up headers.

[mravishankar]

Nice! Wondering if these could be just added to OpBuilder directly?

[pifon2a]

In D138586#3947103, @mravishankar wrote:

Nice! Wondering if these could be just added to OpBuilder directly?

I guess it's a bit harder: there are already clone and cloneWithoutRegions methods that are doing similar things, but not exactly. Also I tried to use BlockAndValueMapping and b.clone(op, bvm), but most of the tests broke.

[mravishankar]

In D138586#3947104, @pifon2a wrote:

In D138586#3947103, @mravishankar wrote:

Nice! Wondering if these could be just added to OpBuilder directly?

I guess it's a bit harder: there are already clone and cloneWithoutRegions methods that are doing similar things, but not exactly. Also I tried to use BlockAndValueMapping and b.clone(op, bvm), but most of the tests broke.

Not saying use the existing methods in OpBuilder, but add these as new methods in OpBuilder....

[pifon2a]

In D138586#3947108, @mravishankar wrote:

In D138586#3947104, @pifon2a wrote:

In D138586#3947103, @mravishankar wrote:

Nice! Wondering if these could be just added to OpBuilder directly?

I guess it's a bit harder: there are already clone and cloneWithoutRegions methods that are doing similar things, but not exactly. Also I tried to use BlockAndValueMapping and b.clone(op, bvm), but most of the tests broke.

Not saying use the existing methods in OpBuilder, but add these as new methods in OpBuilder....

Yes, but wouldn't it be confusing to the users? b.clone(op, newResultsTypes, newOperands) is different from b.clone(op, bvm), where bvm is constructed by mapping old operands to newOperands? If it is not confusing, i can move it there.

[mravishankar]

I dont have an opinion. Maybe River or Mehdi have an opinion on it. Either way, this is also fine for me.

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

bool hasBufferSemantics() {
           .hasBufferSemantics();
     }
 
     bool hasTensorSemantics() {
       return cast<DestinationStyleOpInterface>(*this->getOperation())
           .hasTensorSemantics();
     }
 
-    Operation *clone(OpBuilder & b, Location loc, TypeRange resultTypes,
-                     ValueRange operands) {
-      return cast<DestinationStyleOpInterface>(*this->getOperation())
-          .clone(b, loc, resultTypes, operands);
-    }
-
-    Operation *cloneWithoutRegions(OpBuilder & b, Location loc,
-                                   TypeRange resultTypes, ValueRange operands) {
-      return cast<DestinationStyleOpInterface>(*this->getOperation())
-          .cloneWithoutRegions(b, loc, resultTypes, operands);
-    }
-
     //========================================================================//
     // Helper functions to mutate the `operand_segment_sizes` attribute.
     // These are useful when cloning and changing operand types.
     //========================================================================//
     void setNumInputs(unsigned num) { setOperandSegmentAt(0, num); }
     void setNumOutputBuffers(unsigned num) { setOperandSegmentAt(1, num); }
 
     private:

mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h

 #include "mlir/Support/LLVM.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/Utils/DialectUtilsEnums.h.inc"
 
 namespace mlir {
 
 class OpBuilder;
+class TypeRange;
+class ValueRange;
 
 /// Tests whether the given maps describe a row major matmul. The test is
 /// permutation-invariant. Note that this only checks the affine maps from an
 /// operation, so does not perform any checks on the math being performed within
 /// the reduction.
 bool isRowMajorMatmul(ArrayAttr indexingMaps);
 
 /// Tests whether the given maps describe a column major matmul. The test is
protected:
   OpBuilder &builder;
   MLIRContext *ctx;
   Location loc;
   SmallVector<IteratorTypeT> iterators;
   SmallVector<AffineMap, 4> maps;
   Operation *op;
 };
 
+// Clone the current operation with the operands. This is used to abstract away
+// the optional underlying region creation.
+Operation *clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
+                 ValueRange newOperands);
+
+// Clone the current operation with the operands but leave the regions empty.
+Operation *cloneWithoutRegions(OpBuilder &b, Operation *op,
+                               TypeRange newResultTypes,
+                               ValueRange newOperands);
+
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STRUCTUREDOPSUTILS_H

mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td

InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return llvm::all_of($_op->getOpOperands(),
           [&](OpOperand &opOperand) {
             return isScalar(&opOperand) ||
                    opOperand.get().getType().template isa<RankedTensorType>();
           });
       }]
-    >,
-    //===------------------------------------------------------------------===//
-    // Other static interface methods.
-    //===------------------------------------------------------------------===//
-    InterfaceMethod<
-      /*desc=*/[{
-        Clone the current operation with the given location and operands. This
-        is used to abstract away the optional underlying region creation. This
-        does not change the balance between input, init_buffer and
-        init_tensors operands.
-      }],
-      /*retTy=*/"Operation *",
-      /*methodName=*/"clone",
-      (ins "OpBuilder &":$b, "Location":$loc, "TypeRange":$resultTypes,
-           "ValueRange":$operands),
-      [{
-        BlockAndValueMapping bvm;
-        OperationState state(
-          loc, ConcreteOp::getOperationName(), operands, resultTypes,
-          $_op->getAttrs());
-        for (Region &r : $_op->getRegions())
-          r.cloneInto(state.addRegion(), bvm);
-        return b.create(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, init_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.create(state);
-      }]
     >
   ];
 
   let verify = [{ return detail::verifyDestinationStyleOpInterface($_op); }];
   let verifyWithRegions = 1;
 }
 
 
 #endif // MLIR_DESTINATIONSTYLEOPINTERFACE

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

LogicalResult matchAndRewrite(LinalgOp op,
     for (auto *output : op.getDpsInitOperands()) {
       auto tensorCastOp = output->get().getDefiningOp<tensor::CastOp>();
       bool fold = canFoldIntoConsumerOp(tensorCastOp);
       newOperands.push_back(fold ? tensorCastOp.getOperand() : output->get());
       if (!newOperands.back().getType().isa<MemRefType>())
         newResultTypes.push_back(newOperands.back().getType());
     }
     // Clone op.
-    Operation *newOp =
-        op.clone(rewriter, op->getLoc(), newResultTypes, newOperands);
+    Operation *newOp = clone(rewriter, op, newResultTypes, newOperands);
     SmallVector<Value, 4> replacements;
     replacements.reserve(newOp->getNumResults());
     for (auto result : llvm::zip(op->getResults(), newOp->getResults())) {
       Value oldResult = std::get<0>(result);
       Value newResult = std::get<1>(result);
       if (newResult.getType() != oldResult.getType()) {
         replacements.push_back(rewriter.create<tensor::CastOp>(
             op->getLoc(), oldResult.getType(), newResult));
LogicalResult matchAndRewrite(tensor::CastOp castOp,
     SmallVector<Value> newOperands{linalgOp.getDpsInputOperands()};
     SmallVector<Value> outputOperands{linalgOp.getDpsInitOperands()};
     outputOperands[resultNumber] = newOperand;
     newOperands.append(outputOperands.begin(), outputOperands.end());
 
     SmallVector<Type> resultTypes(linalgOp->result_type_begin(),
                                   linalgOp->result_type_end());
     resultTypes[resultNumber] = resultType;
-    Operation *newOp = linalgOp.clone(rewriter, loc, resultTypes, newOperands);
+    Operation *newOp = clone(rewriter, linalgOp, resultTypes, newOperands);
 
     // Create a tensor.cast operation back to the original type.
     Value castBack = rewriter.create<tensor::CastOp>(
         loc, resultValue.getType(), newOp->getResult(resultNumber));
 
     SmallVector<Value> results(newOp->result_begin(), newOp->result_end());
     results[resultNumber] = castBack;
     rewriter.replaceOp(linalgOp, results);
LogicalResult matchAndRewrite(LinalgOp linalgOp,
     }
 
     // If the generic op has all the required static information, no
     // canonicalization needed.
     if (!changeNeeded)
       return failure();
 
     // Clone op.
-    Operation *newOp =
-        linalgOp.clone(rewriter, linalgOp->getLoc(), resultTypes, newOperands);
+    Operation *newOp = clone(rewriter, linalgOp, resultTypes, newOperands);
     SmallVector<Value> replacements;
     replacements.reserve(newOp->getNumResults());
     for (auto it : llvm::zip(linalgOp->getResults(), newOp->getResults())) {
       Value newResult = std::get<1>(it);
       Value oldResult = std::get<0>(it);
       Type newType = newResult.getType();
       Type oldType = oldResult.getType();
       replacements.push_back(

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

SmallVector<Value> tiledOperands =
                         /*omitPartialTileCheck=*/true);
 
     SmallVector<Type, 4> resultTensorTypes;
     for (OpOperand *opOperand : linalgOp.getDpsInitOperands())
       resultTensorTypes.push_back(
           tiledOperands[opOperand->getOperandNumber()].getType());
 
     Operation *newOp =
-        linalgOp.clone(rewriter, linalgLoc, resultTensorTypes, tiledOperands);
+        clone(rewriter, linalgOp, resultTensorTypes, tiledOperands);
     rewriter.replaceOp(sliceOp, newOp->getResults());
     return success();
   }
 };
 } // namespace
 
 void mlir::linalg::populateBubbleUpExtractSliceOpPatterns(
     RewritePatternSet &patterns) {
   auto *context = patterns.getContext();
   patterns.add<BubbleUpExtractSliceOpPattern>(context);
 }

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

bufferizeDestinationStyleOpInterface(RewriterBase &rewriter,
   newOperands.append(newOutputBuffers.begin(), newOutputBuffers.end());
 
   // Set insertion point now that potential alloc/dealloc are introduced.
   rewriter.setInsertionPoint(op);
   // Clone the op, but use the new operands. Move the existing block into the
   // new op. Since the new op does not have any tensor results, it does not
   // return anything.
   assert(op->getNumRegions() == 1 && "expected that op has 1 region");
-  auto newOp = cast<DestinationStyleOpInterface>(op.cloneWithoutRegions(
-      rewriter, op.getLoc(), /*resultTypes=*/TypeRange{}, newOperands));
+  auto newOp = cast<DestinationStyleOpInterface>(cloneWithoutRegions(
+      rewriter, op, /*resultTypes=*/TypeRange{}, newOperands));
   rewriter.inlineRegionBefore(op->getRegion(0), newOp->getRegion(0),
                               newOp->getRegion(0).begin());
 
   // Replace the results of the old op with the new output buffers.
   replaceOpWithBufferizedValues(rewriter, op, newOutputBuffers);
 
   return success();
 }

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

for (OpOperand *operand : producer.getDpsInitOperands()) {
     SmallVector<int64_t, 4> staticSizesVector(rank, ShapedType::kDynamic);
     SmallVector<int64_t, 4> staticStridesVector(
         rank, ShapedType::kDynamic);
     resultTypes.push_back(tensor::ExtractSliceOp::inferResultType(
         tensorType, staticOffsetsVector, staticSizesVector,
         staticStridesVector));
   }
 
-  Operation *clonedOp = producer.clone(b, loc, resultTypes, clonedShapes);
+  Operation *clonedOp = clone(b, producer, resultTypes, clonedShapes);
 
   // Shift all IndexOp results by the tile offset.
   SmallVector<OpFoldResult> allIvs = llvm::to_vector(
       llvm::map_range(loopRanges, [&](Range range) { return range.offset; }));
   offsetIndices(b, clonedOp, allIvs);
 
   return clonedOp;
 }

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

if (iterArg) {
     iterArg->set(outputOperand->get());
     tiledOperands[outputOperand->getOperandNumber()] = sliceOp.getResult();
   }
 
   // Clone the producer using the tiled producer operands.
   TypeRange resultTypes = ValueRange(tiledOperands)
                               .take_back(producerOp.getNumDpsInits())
                               .getTypes();
-  LinalgOp clonedOp = producerOp.clone(b, loc, resultTypes, tiledOperands);
+  LinalgOp clonedOp = clone(b, producerOp, resultTypes, tiledOperands);
 
   // Shift all IndexOp results by the tile offset.
   offsetIndices(b, clonedOp, allIvs);
 
   return clonedOp;
 }
 
 //===----------------------------------------------------------------------===//

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

SmallVector<OpFoldResult> sizeBounds =
                                                  allShapeSizes);
     SmallVector<Value> tiledOperands = makeTiledShapes(
         b, loc, op, valuesToTile, getAsOpFoldResult(interchangedIvs), tileSizes,
         sizeBounds,
         /*omitPartialTileCheck=*/false);
 
     SmallVector<Type> resultTensorTypes =
         getTensorOutputTypes(op, tiledOperands);
-    res = op.clone(b, loc, resultTensorTypes, tiledOperands);
+    res = clone(b, op, resultTensorTypes, tiledOperands);
     tensorResults =
         insertSlicesBack(builder, loc, op, tiledOperands, res->getResults());
     return scf::ValueVector(tensorResults.begin(), tensorResults.end());
   };
   GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
                                  tiledLoopBodyBuilder, procInfo);
 
   // 3. Transform IndexOp results w.r.t. the tiling.

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

getTiledImplementation(Operation *op, OpBuilder &b,
     LinalgOp linalgOp = cast<LinalgOp>(op);
     SmallVector<Value> valuesToTile = linalgOp->getOperands();
     SmallVector<Value, 4> tiledOperands = makeTiledShapes(
         b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
 
     SmallVector<Type> resultTensorTypes =
         getTensorOutputTypes(linalgOp, tiledOperands);
 
-    Operation *tiledOp =
-        linalgOp.clone(b, loc, resultTensorTypes, tiledOperands);
+    Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
     offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
 
     return {tiledOp};
   }
 
   // Return the details of the output tile generated by the tiled
   // implementation.
   LogicalResult

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

if (failed(cast<ReifyRankedShapedTypeOpInterface>(opToPad.getOperation())
                  .reifyResultShapes(b, reifiedResultShapes)))
     return failure();
   assert(reifiedResultShapes.size() == opToPad->getNumResults() &&
          "expected same number of results");
 
   // Clone `opToPad` to operate on the statically padded shapes.
   auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumDpsInits()).getTypes();
-  paddedOp = opToPad.clone(b, loc, resultTensorTypes, newOperands);
+  paddedOp = clone(b, opToPad, resultTensorTypes, newOperands);
 
   // Recover the slice out of the new static results. This keeps the original
   // linalg op around because it uses the dims of the original results.
   SmallVector<Value> paddedSubviewResults;
   paddedSubviewResults.reserve(opToPad->getNumResults());
   for (const auto &en : llvm::enumerate(paddedOp->getResults())) {
     Value paddedResult = en.value();
     int64_t resultNumber = en.index();

mlir/lib/Dialect/Utils/StructuredOpsUtils.cpp

 //===- StructuredOpsUtils.cpp - Utilities used by structured 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/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 
 #include "mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc"
 
 using namespace mlir;
 
 bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
   if (indexingMaps.size() != 3)
bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
   AffineExpr k = map0.getResult(2);
   auto *context = indexingMaps.getContext();
   auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
   auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
   auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
   auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
   return indexingMaps == maps;
 }
+
+Operation *mlir::clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
+                       ValueRange newOperands) {
+  BlockAndValueMapping bvm;
+  OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
+                       op->getAttrs());
+  for (Region &r : op->getRegions())
+    r.cloneInto(state.addRegion(), bvm);
+  return b.create(state);
+}
+
+Operation *mlir::cloneWithoutRegions(OpBuilder &b, Operation *op,
+                                     TypeRange newResultTypes,
+                                     ValueRange newOperands) {
+  OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
+                       op->getAttrs());
+  for (size_t cnt = 0, e = op->getNumRegions(); cnt < e; ++cnt)
+    state.addRegion();
+  return b.create(state);
+}