From 498c685b46a7faf1b48ea014aa1755e95c870319 Mon Sep 17 00:00:00 2001
From: Philippe Tillet <phil@openai.com>
Date: Fri, 7 Oct 2022 21:52:29 -0700
Subject: [PATCH] [OPTIMIZER] layout simplification: ignore non-tensor iter
 arguments in for loop rematerialization (#749)

---
 lib/Dialect/TritonGPU/Transforms/Combine.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/lib/Dialect/TritonGPU/Transforms/Combine.cpp b/lib/Dialect/TritonGPU/Transforms/Combine.cpp
index a354993a8..5574b3ffd 100644
--- 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());