[OPTIMIZER] layout simplification: ignore non-tensor iter arguments in for loop rematerialization (#749)

2022-10-07 21:52:29 -07:00
parent e843257295
commit 498c685b46
1 changed files with 2 additions and 0 deletions
--- a/lib/Dialect/TritonGPU/Transforms/Combine.cpp
+++ b/lib/Dialect/TritonGPU/Transforms/Combine.cpp
@@ -306,6 +306,8 @@ public:
        auto fwdCvtIt = std::find_if(opIt, fwdEndIt, isCvt);
        auto bwdCvtIt = std::find_if(bwdBeginIt, opIt, isCvt);

+        if (!iterArg.value().getType().isa<RankedTensorType>())
+          continue;
        if (fwdCvtIt != fwdEndIt) {
          auto newFor = tryConvertIterArg(forOp, rewriter, iterArg.index(),
                                          (*fwdCvtIt)->getResult(0).getType());