[mlir] Use getArgNumber to get argument index (NFC)

[linuxlonelyeagle]

This PR simplifies the process of obtaining the argument index.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: lonely eagle (linuxlonelyeagle)

Changes

This PR simplifies the process of obtaining the argument index.

---

Full diff: https://github.com/llvm/llvm-project/pull/162959.diff

1 Files Affected:

• (modified) mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp (+1-2)

diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 111f58ef92f51..b6ca18c760ebf 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -324,8 +324,7 @@ static LoopLikeOpInterface hoistSubsetAtIterArg(RewriterBase &rewriter,
                                                 LoopLikeOpInterface loopLike,
                                                 BlockArgument iterArg) {
   assert(iterArg.getOwner()->getParentOp() == loopLike && "invalid iter_arg");
-  BlockArgument *it = llvm::find(loopLike.getRegionIterArgs(), iterArg);
-  int64_t iterArgIdx = std::distance(loopLike.getRegionIterArgs().begin(), it);
+  int64_t iterArgIdx = iterArg.getArgNumber() - 1;
   MatchingSubsets subsets;
   if (failed(subsets.populateSubsetOpsAtIterArg(loopLike, iterArg)))
     return loopLike;

[PragmaTwice]

Hmm I think it can be iterArg.getArgNumber() - N depending on how getRegionIterArgs() is defined for the specific looplike op? For scf.for it can always be N = 1, but is it applicable for other ops (forall, while..)?

[linuxlonelyeagle]

Thank you for pointing this out. However, there's a fact that the tests for loop-invariant-subset-hoisting only include scf.for. Is this pass designed solely for scf.for loops?