[mlir][linalg] Cleanup LinalgOp usage in promotion.

[gysit]

Summary

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

[stella.stamenova]

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

335

Since operandNumber is already an int64_t, this cast is unnecessary.

[stella.stamenova]

I would not consider this an NFC and I think the title should be updated to not include the NFC designation. It might not be changing the intended functionality, but it is changing the logic (and functions called) to achieve the result.

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

327

I thought getOperandNumber returned an unsigned? I would expect assigning an unsigned value to an int64_t to generate a warning.

[gysit]

Address comment.

[gysit]

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

327

I will check if there is an error but I believe the cast happens implicitly here:

c.f. "The compiler doesn't warn about implicit conversions between signed and unsigned integral types." taken from https://docs.microsoft.com/en-us/cpp/cpp/type-conversions-and-type-safety-modern-cpp?view=msvc-160

335

Thanks for the hint!

[stella.stamenova]

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

327

The Visual Studio compiler, clang and gcc often behave very differently and raise different warnings. They are also all supported by llvm/mlir/etc. Generally, though, the expectation is that the code will be warning-free when building with clang, so I would recommend trying that to see if there are any warnings.

[gysit]

Rebase.

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

 };
 } // namespace
 
 LinalgOpInstancePromotionOptions::LinalgOpInstancePromotionOptions(
     LinalgOp linalgOp, const LinalgPromotionOptions &options)
     : subViews(), dynamicBuffers(options.dynamicBuffers),
       alignment(options.alignment) {
   assert(linalgOp.hasBufferSemantics() && "revisit usage of shaped operand");
-  int64_t nBuffers = linalgOp.getNumShapedOperands();
   auto vUseFullTileBuffers =
       options.useFullTileBuffers.getValueOr(llvm::SmallBitVector());
-  vUseFullTileBuffers.resize(nBuffers, options.useFullTileBuffersDefault);
+  vUseFullTileBuffers.resize(linalgOp.getNumInputsAndOutputs(),
+                             options.useFullTileBuffersDefault);
 
-  for (int64_t idx = 0; idx != nBuffers; ++idx) {
-    if (options.operandsToPromote && !options.operandsToPromote->count(idx))
+  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+    int64_t operandNumber = opOperand->getOperandNumber();
+    if (options.operandsToPromote &&
+        !options.operandsToPromote->count(operandNumber))
       continue;
-    auto *op = linalgOp.getShapedOperand(idx).getDefiningOp();
+    Operation *op = opOperand->get().getDefiningOp();
     if (auto sv = dyn_cast_or_null<memref::SubViewOp>(op)) {
-      subViews[idx] = sv;
-      useFullTileBuffers[sv] = vUseFullTileBuffers[idx];
+      subViews[operandNumber] = sv;
+      useFullTileBuffers[sv] = vUseFullTileBuffers[operandNumber];
     }
   }
 
   if (options.allocationFn) {
     allocationFn = *options.allocationFn;
   } else {
     allocationFn = [&](OpBuilder &b, memref::SubViewOp subViewOp,
                        ArrayRef<Value> boundingSubViewSize,
promoteSubViews(ImplicitLocOpBuilder &b, LinalgOp op,
   if (!promotedBuffersAndViews ||
       promotedBuffersAndViews->size() != options.subViews.size())
     return {};
 
   // 2. Append all other operands as they appear, this enforces that such
   // operands are not views. This is to support cases such as FillOp taking
   // extra scalars etc.  Keep a reference to output buffers;
   SmallVector<Value, 8> opViews;
-  opViews.reserve(op.getNumShapedOperands());
+  opViews.reserve(op.getNumInputsAndOutputs());
   SmallVector<std::pair<Value, Value>, 8> writebackViews;
   writebackViews.reserve(promotedBuffersAndViews->size());
-  for (auto view : llvm::enumerate(op.getShapedOperands())) {
-    if (options.subViews.count(view.index()) != 0) {
-      if (options.useFullTileBuffers[view.value()])
+  for (OpOperand *opOperand : op.getInputAndOutputOperands()) {
+    int64_t operandNumber = opOperand->getOperandNumber();
+    if (options.subViews.count(operandNumber) != 0) {
+      if (options.useFullTileBuffers[opOperand->get()])
         opViews.push_back(
-            (*promotedBuffersAndViews)[view.index()].fullLocalView);
+            (*promotedBuffersAndViews)[operandNumber].fullLocalView);
       else
         opViews.push_back(
-            (*promotedBuffersAndViews)[view.index()].partialLocalView);
-      if (static_cast<int64_t>(view.index()) >= op.getNumInputs())
+            (*promotedBuffersAndViews)[operandNumber].partialLocalView);
+      if (operandNumber >= op.getNumInputs())
         writebackViews.emplace_back(std::make_pair(
-            view.value(),
-            (*promotedBuffersAndViews)[view.index()].partialLocalView));
+            opOperand->get(),
+            (*promotedBuffersAndViews)[operandNumber].partialLocalView));
     } else {
-      opViews.push_back(view.value());
+      opViews.push_back(opOperand->get());
     }
   }
   op->setOperands(0, opViews.size(), opViews);
 
   OpBuilder::InsertionGuard guard(b);
   b.setInsertionPointAfter(op);
   // 3. Emit write-back for the promoted output views: copy the partial view.
   for (auto viewAndPartialLocalView : writebackViews) {
     if (failed(options.copyOutFn(b, viewAndPartialLocalView.second,
                                  viewAndPartialLocalView.first)))
       return {};
   }
 
   // 4. Dealloc all local buffers.
   for (const auto &pi : *promotedBuffersAndViews)
     (void)options.deallocationFn(b, pi.second.fullLocalView);
   return op;
 }
 
 LogicalResult
 mlir::linalg::promoteSubviewsPrecondition(Operation *op,
                                           LinalgPromotionOptions options) {
-  LinalgOp linOp = dyn_cast<LinalgOp>(op);
+  LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   // Transformation applies to buffers only.
-  if (!linOp || !linOp.hasBufferSemantics())
+  if (!linalgOp || !linalgOp.hasBufferSemantics())
     return failure();
   // Check that at least one of the requested operands is indeed a subview.
-  for (auto en : llvm::enumerate(linOp.getShapedOperands())) {
-    auto sv = isa_and_nonnull<memref::SubViewOp>(en.value().getDefiningOp());
+  for (OpOperand *opOperand : linalgOp.getInputAndOutputOperands()) {
+    auto sv =
+        isa_and_nonnull<memref::SubViewOp>(opOperand->get().getDefiningOp());
     if (sv) {
       if (!options.operandsToPromote.hasValue() ||
-          options.operandsToPromote->count(en.index()))
+          options.operandsToPromote->count(opOperand->getOperandNumber()))
         return success();
     }
   }
   // TODO: Check all subviews requested are bound by a static constant.
   // TODO: Check that the total footprint fits within a given size.
   return failure();
 }