From e948a618b30ef32ac6efb41b0c1e5740dd5f70fd Mon Sep 17 00:00:00 2001
From: goostavz <109190422+goostavz@users.noreply.github.com>
Date: Mon, 17 Oct 2022 11:29:55 +0800
Subject: [PATCH] [Triton-MLIR] fix a tiny bug in coalesce pass (#782)

---
 lib/Dialect/TritonGPU/Transforms/Coalesce.cpp | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp b/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp
index d00c00c2d..019a91d01 100644
--- a/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp
+++ b/lib/Dialect/TritonGPU/Transforms/Coalesce.cpp
@@ -63,9 +63,13 @@ struct CoalescePass : public TritonGPUCoalesceBase<CoalescePass> {
     auto convertType = getTypeConverter(axisInfo, ptr, numWarps);
     // convert operands
     SmallVector<Value, 4> newArgs;
-    for (auto v : op->getOperands())
-      newArgs.push_back(builder.create<triton::gpu::ConvertLayoutOp>(
-          op->getLoc(), convertType(v.getType()), v));
+    for (auto v : op->getOperands()) {
+      if (v.getType().isa<RankedTensorType>())
+        newArgs.push_back(builder.create<triton::gpu::ConvertLayoutOp>(
+            op->getLoc(), convertType(v.getType()), v));
+      else
+        newArgs.push_back(v);
+    }
     // convert output types
     SmallVector<Type, 4> newTypes;
     for (auto t : op->getResultTypes()) {