special encoding for broadcast

2022-06-18 21:16:45 +08:00
parent 53cf93ce6a
commit 9d1b5e3f79
6 changed files with 248 additions and 72 deletions
--- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPU.cpp
+++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPU.cpp
@@ -342,9 +342,13 @@ public:
    //    mlir::scf::populateSCFStructurealTypeConversionsAndLegality(...) here?
    populateSCFPatterns(typeConverter, patterns);

-    if(failed(applyPartialConversion(mod, target, 
-                                      std::move(patterns))))
+    if(failed(applyPartialConversion(mod, target, std::move(patterns))))
        return signalPassFailure();
+
+    // update layouts
+    //  broadcast src => multicast, dst => broadcasted
+    if(failed(target.refineLayouts(mod, numWarps)))
+      return signalPassFailure();
  }
 };