Move getMakeTensorPtrOp func to Triton/IR/Utility.h (#6755)

anmyachev · web-flow · commit 3c8cdbfb611f · 2025-05-08T15:19:06.000-07:00
To fix the problem:
```bash
  In file included from /runner/_work/intel-xpu-backend-for-triton/intel-xpu-backend-for-triton/include/triton/Dialect/TritonGPU/IR/Dialect.h:12,
                   from /runner/_work/intel-xpu-backend-for-triton/intel-xpu-backend-for-triton/include/triton/Analysis/Utility.h:8,
                   from /runner/_work/intel-xpu-backend-for-triton/intel-xpu-backend-for-triton/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp:9:
  /runner/_work/intel-xpu-backend-for-triton/intel-xpu-backend-for-triton/include/triton/Dialect/TritonGPU/IR/Types.h:12:10: fatal error: triton/Dialect/TritonGPU/IR/TypeInterfaces.h.inc: No such file or directory
     12 | #include "triton/Dialect/TritonGPU/IR/TypeInterfaces.h.inc"
        |          ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  compilation terminated.
```

Signed-off-by: Anatoly Myachev &lt;anatoly.myachev@intel.com&gt;
diff --git a/include/triton/Analysis/Utility.h b/include/triton/Analysis/Utility.h
@@ -411,8 +411,6 @@ template <typename T> class CallGraph {
 // Create a basic DataFlowSolver with constant and dead code analysis included.
 std::unique_ptr<DataFlowSolver> createDataFlowSolver();
 
-triton::MakeTensorPtrOp getMakeTensorPtrOp(Value v);
-
 } // namespace mlir
 
 #endif // TRITON_ANALYSIS_UTILITY_H
diff --git a/include/triton/Dialect/Triton/IR/Utility.h b/include/triton/Dialect/Triton/IR/Utility.h
@@ -177,6 +177,8 @@ template <typename T> auto seq(T start, T end, T step) {
 Value getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
                   Value pred);
 
+MakeTensorPtrOp getMakeTensorPtrOp(Value v);
+
 } // namespace triton
 } // namespace mlir
 
diff --git a/lib/Analysis/Utility.cpp b/lib/Analysis/Utility.cpp
@@ -1029,75 +1029,4 @@ std::unique_ptr<DataFlowSolver> createDataFlowSolver() {
   return solver;
 }
 
-static MakeTensorPtrOp getMakeTensorPtrOpImpl(Operation *op, Value v) {
-
-  if (auto makeTensorPtrOp = dyn_cast<MakeTensorPtrOp>(op)) {
-    return makeTensorPtrOp;
-  }
-
-  if (auto advanceOp = dyn_cast<AdvanceOp>(op)) {
-    return getMakeTensorPtrOp(advanceOp.getPtr());
-  }
-
-  if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
-    auto idx = cast<OpResult>(v).getResultNumber();
-    llvm::SmallVector<scf::YieldOp> yieldOps;
-    op->walk([&](Operation *op) {
-      if (auto yieldOp = dyn_cast<scf::YieldOp>(op))
-        yieldOps.push_back(yieldOp);
-    });
-
-    // benzh@ if multi yields, all yields operand should come from same arg.
-    Value newValue = yieldOps[0].getOperands()[idx];
-    return getMakeTensorPtrOp(newValue);
-  }
-
-  llvm_unreachable("Unable to getMakeTensorPtr()");
-}
-
-MakeTensorPtrOp getMakeTensorPtrOp(Value v) {
-  using BranchOps = llvm::SetVector<std::pair<Operation *, int>>;
-  llvm::DenseMap<Block *, BranchOps> blockToCFOps;
-  auto moduleOp =
-      v.getParentBlock()->getParentOp()->getParentOfType<ModuleOp>();
-
-  moduleOp.walk([&](Operation *op) {
-    if (auto br = dyn_cast<cf::BranchOp>(op)) {
-      Block *block = br.getDest();
-      blockToCFOps[block].insert({op, -1});
-    }
-    if (auto condBr = dyn_cast<cf::CondBranchOp>(op)) {
-      Block *blockT = condBr.getTrueDest();
-      Block *blockF = condBr.getFalseDest();
-      blockToCFOps[blockT].insert({condBr, 1});
-      blockToCFOps[blockF].insert({condBr, 0});
-    }
-  });
-
-  if (Operation *definingOp = v.getDefiningOp())
-    return getMakeTensorPtrOpImpl(definingOp, v);
-
-  // If there is no defining op, v must be a BlockArgument.
-  BlockArgument arg = cast<BlockArgument>(v);
-  unsigned argNum = arg.getArgNumber();
-  Operation *argOwner = arg.getOwner()->getParentOp();
-
-  if (auto forOp = dyn_cast<scf::ForOp>(argOwner))
-    return getMakeTensorPtrOp(
-        forOp.getOperand(argNum + forOp.getNumControlOperands() - 1));
-  if (auto funcOp = dyn_cast<FunctionOpInterface>(argOwner)) {
-    Block *block = arg.getOwner();
-    Operation *op;
-    int tOrF;
-    std::tie(op, tOrF) = blockToCFOps[block][0];
-    if (auto br = dyn_cast<cf::BranchOp>(op))
-      return getMakeTensorPtrOp(br.getDestOperands()[argNum]);
-    if (auto condBr = dyn_cast<cf::CondBranchOp>(op))
-      return getMakeTensorPtrOp(tOrF ? condBr.getTrueDestOperands()[argNum]
-                                     : condBr.getFalseDestOperands()[argNum]);
-    return getMakeTensorPtrOp(argOwner->getOperand(argNum));
-  }
-  llvm_unreachable("Unable to getMakeTensorPtr()");
-}
-
 } // namespace mlir
diff --git a/lib/Dialect/Triton/IR/Utility.cpp b/lib/Dialect/Triton/IR/Utility.cpp
@@ -1,4 +1,5 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 using namespace mlir;
@@ -17,3 +18,75 @@ Value tt::getPredMask(RewriterBase &rewriter, Type typeLike, Value currentMask,
   }
   return mask;
 }
+
+static tt::MakeTensorPtrOp getMakeTensorPtrOpImpl(Operation *op, Value v) {
+
+  if (auto makeTensorPtrOp = dyn_cast<tt::MakeTensorPtrOp>(op)) {
+    return makeTensorPtrOp;
+  }
+
+  if (auto advanceOp = dyn_cast<tt::AdvanceOp>(op)) {
+    return tt::getMakeTensorPtrOp(advanceOp.getPtr());
+  }
+
+  if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
+    auto idx = cast<OpResult>(v).getResultNumber();
+    llvm::SmallVector<scf::YieldOp> yieldOps;
+    op->walk([&](Operation *op) {
+      if (auto yieldOp = dyn_cast<scf::YieldOp>(op))
+        yieldOps.push_back(yieldOp);
+    });
+
+    // benzh@ if multi yields, all yields operand should come from same arg.
+    Value newValue = yieldOps[0].getOperands()[idx];
+    return tt::getMakeTensorPtrOp(newValue);
+  }
+
+  llvm_unreachable("Unable to getMakeTensorPtr()");
+}
+
+tt::MakeTensorPtrOp tt::getMakeTensorPtrOp(Value v) {
+  using BranchOps = llvm::SetVector<std::pair<Operation *, int>>;
+  llvm::DenseMap<Block *, BranchOps> blockToCFOps;
+  auto moduleOp =
+      v.getParentBlock()->getParentOp()->getParentOfType<ModuleOp>();
+
+  moduleOp.walk([&](Operation *op) {
+    if (auto br = dyn_cast<cf::BranchOp>(op)) {
+      Block *block = br.getDest();
+      blockToCFOps[block].insert({op, -1});
+    }
+    if (auto condBr = dyn_cast<cf::CondBranchOp>(op)) {
+      Block *blockT = condBr.getTrueDest();
+      Block *blockF = condBr.getFalseDest();
+      blockToCFOps[blockT].insert({condBr, 1});
+      blockToCFOps[blockF].insert({condBr, 0});
+    }
+  });
+
+  if (Operation *definingOp = v.getDefiningOp())
+    return getMakeTensorPtrOpImpl(definingOp, v);
+
+  // If there is no defining op, v must be a BlockArgument.
+  BlockArgument arg = cast<BlockArgument>(v);
+  unsigned argNum = arg.getArgNumber();
+  Operation *argOwner = arg.getOwner()->getParentOp();
+
+  if (auto forOp = dyn_cast<scf::ForOp>(argOwner))
+    return tt::getMakeTensorPtrOp(
+        forOp.getOperand(argNum + forOp.getNumControlOperands() - 1));
+  if (auto funcOp = dyn_cast<FunctionOpInterface>(argOwner)) {
+    Block *block = arg.getOwner();
+    Operation *op;
+    int tOrF;
+    std::tie(op, tOrF) = blockToCFOps[block][0];
+    if (auto br = dyn_cast<cf::BranchOp>(op))
+      return tt::getMakeTensorPtrOp(br.getDestOperands()[argNum]);
+    if (auto condBr = dyn_cast<cf::CondBranchOp>(op))
+      return tt::getMakeTensorPtrOp(
+          tOrF ? condBr.getTrueDestOperands()[argNum]
+               : condBr.getFalseDestOperands()[argNum]);
+    return tt::getMakeTensorPtrOp(argOwner->getOperand(argNum));
+  }
+  llvm_unreachable("Unable to getMakeTensorPtr()");
+}
diff --git a/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp
@@ -6,8 +6,8 @@
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
-#include "triton/Analysis/Utility.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 
 using namespace mlir;
@@ -340,7 +340,7 @@ class RewriteTensorPointerPass
         continue;
       }
       needRewrite = true;
-      auto makeTensorPtrOp = getMakeTensorPtrOp(results[i]);
+      auto makeTensorPtrOp = triton::getMakeTensorPtrOp(results[i]);
       assert(rewritedInfo.count(makeTensorPtrOp.getResult()));
       const auto &info = rewritedInfo[makeTensorPtrOp.getResult()];
       for (unsigned j = 0; j < info.length(); ++j) {
@@ -378,7 +378,7 @@ class RewriteTensorPointerPass
         oldResIdx++;
         newResIdx++;
       } else {
-        auto makeTensorPtrOp = getMakeTensorPtrOp(results[oldResIdx]);
+        auto makeTensorPtrOp = triton::getMakeTensorPtrOp(results[oldResIdx]);
         assert(rewritedInfo.count(makeTensorPtrOp.getResult()));
         auto info = rewritedInfo[makeTensorPtrOp.getResult()];
         for (unsigned j = 0; j < info.length(); ++j) {
