[Mlir-commits] [mlir] f44e90b - [mlir][linalg] Cleanup LinalgOp usage in scalar inlining.

[Tobias Gysi]

Author: Tobias Gysi
Date: 2021-06-03T14:45:14Z
New Revision: f44e90b93ad89d0a93a550e8fca8a2a9b6e4f190

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

LOG: [mlir][linalg] Cleanup LinalgOp usage in scalar inlining.

Replace the uses of deprecated Structured Op Interface methods in InlineScalarOperands.cpp. This patch is based on https://reviews.llvm.org/D103394.

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp b/mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp
index aa01029471ce..ed0b1b5a40ac 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp
@@ -34,26 +34,27 @@ struct InlineScalarOperands : public OpRewritePattern<GenericOp> {
     SmallVector<size_t> scalarOperands;
     SmallVector<AffineMap> newIndexingMaps;
     SmallVector<Value> newOperands;
-    for (auto it : llvm::enumerate(llvm::zip(genericOp.getInputIndexingMaps(),
-                                             genericOp.getInputTensors()))) {
-      AffineMap map = std::get<0>(it.value());
-      if (map.isConstant()) {
-        scalarOperands.emplace_back(it.index());
+    for (OpOperand *opOperand : genericOp.getInputOperands()) {
+      AffineMap map = genericOp.getTiedIndexingMap(opOperand);
+      if (genericOp.isInputTensor(opOperand) && map.isConstant()) {
+        scalarOperands.emplace_back(opOperand->getOperandNumber());
       } else {
         newIndexingMaps.emplace_back(map);
-        newOperands.emplace_back(std::get<1>(it.value()));
+        newOperands.emplace_back(opOperand->get());
       }
     }
 
     if (scalarOperands.empty())
       return failure();
 
-    newIndexingMaps.append(genericOp.getOutputIndexingMaps());
+    for (OpOperand *opOperand : genericOp.getOutputOperands())
+      newIndexingMaps.emplace_back(genericOp.getTiedIndexingMap(opOperand));
 
     Location loc = genericOp->getLoc();
+    SmallVector<Value> outputOperands = genericOp.getOutputOperands();
     auto newOp = rewriter.create<GenericOp>(
-        loc, genericOp->getResultTypes(), newOperands,
-        genericOp.getOutputTensors(), newIndexingMaps,
+        loc, genericOp->getResultTypes(), newOperands, outputOperands,
+        newIndexingMaps,
         llvm::to_vector<4>(
             genericOp.iterator_types().template getAsValueRange<StringAttr>()));
     rewriter.cloneRegionBefore(genericOp.region(), newOp.region(),
@@ -64,14 +65,15 @@ struct InlineScalarOperands : public OpRewritePattern<GenericOp> {
     rewriter.setInsertionPointToStart(body);
 
     for (auto idx : llvm::reverse(scalarOperands)) {
-      Value operand = genericOp.getInput(idx);
-      AffineMap map = genericOp.getInputIndexingMap(idx);
+      OpOperand *opOperand = genericOp.getInputOperand(idx);
+      AffineMap map = genericOp.getTiedIndexingMap(opOperand);
       SmallVector<int64_t> indices = map.getConstantResults();
       SmallVector<Value> indicesValues;
       for (auto idx : indices)
         indicesValues.emplace_back(rewriter.create<ConstantIndexOp>(loc, idx));
-      operand = rewriter.create<tensor::ExtractOp>(loc, operand, indicesValues);
-      body->getArgument(idx).replaceAllUsesWith(operand);
+      Value extractedValue = rewriter.create<tensor::ExtractOp>(
+          loc, opOperand->get(), indicesValues);
+      body->getArgument(idx).replaceAllUsesWith(extractedValue);
       body->eraseArgument(idx);
     }