Enhance the removelayout

intel · May 8, 2024 · 7b22c45 · 7b22c45
1 parent 31f07fb
commit 7b22c45
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp b/third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp
@@ -1093,6 +1093,13 @@ void LayoutRematerialization::hoistConvertOnTopOfExtOrBroadcast() {
 void LayoutRematerialization::backwardRematerialization(
     ConvertLayoutOp convertOp) {
   RankedTensorType targetType = convertOp.getType();
+  // we don't backward propagate the dot layout with blocked layout as parent.
+  // It introduces a lot of duplicated values in multiple-threads.
+  if (auto dotLayout =
+          dyn_cast<DotOperandEncodingAttr>(targetType.getEncoding())) {
+    if (dotLayout.getParent().isa<BlockedEncodingAttr>())
+      return;
+  }
   Value oldV = convertOp->getOperand(0);
   LDBG("check backward remat with source " << oldV << " encoding "
                                            << targetType.getEncoding());