[RemoveLayoutConversions]: Protect for loop support with env. variable (#5282)

Signed-off-by: Ettore Tiotto <[email protected]>

Author: etiotto

include/triton/Tools/Sys/GetEnv.hpp

@@ -50,6 +50,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_INTEL_FAST_MATH",
     "TRITON_INTEL_ONE_MATRIX_PER_LOAD_BT",
     "TRITON_INTEL_PREDICATED",
+    "TRITON_INTEL_REMOVELAYOUTCONVERSION_SUPPORT_FOR_LOOP",
     // clang-format on
 };
 

test/TritonIntelGPU/combine.mlir

@@ -1,4 +1,5 @@
-// RUN: triton-opt %s -split-input-file -allow-unregistered-dialect -tritonintelgpu-remove-layout-conversions 2>&1 | FileCheck %s
+// RUN: env TRITON_INTEL_REMOVELAYOUTCONVERSION_SUPPORT_FOR_LOOP=0 triton-opt %s -split-input-file -allow-unregistered-dialect -tritonintelgpu-remove-layout-conversions 2>&1 | FileCheck --check-prefixes=CHECK %s
+// RUN: env TRITON_INTEL_REMOVELAYOUTCONVERSION_SUPPORT_FOR_LOOP=1 triton-opt %s -split-input-file -allow-unregistered-dialect -tritonintelgpu-remove-layout-conversions 2>&1 | FileCheck --check-prefixes=CHECK,FOR-SUPPORT %s
 
 #layout0 = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
 #layout1 = #ttg.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
@@ -351,7 +352,7 @@ tt.func @loop(%arg0: !tt.ptr<f32>, %arg1: i32, %arg2: !tt.ptr<f32>, %arg3: i32,
 //     CHECK: scf.for
 // CHECK-NOT:   ttg.convert_layout
 //     CHECK:   scf.if
-//     CHECK:     ttg.convert_layout
+//     FOR-SUPPORT: ttg.convert_layout
 //     CHECK:     scf.yield
 // CHECK-NEXT:  else
 // CHECK-NEXT:    scf.yield

third_party/intel/lib/TritonIntelGPUTransforms/RemoveLayoutConversions.cpp

@@ -25,6 +25,7 @@
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
+#include "triton/Tools/Sys/GetEnv.hpp"
 #include "llvm/ADT/TypeSwitch.h"
 #include <deque>
 
@@ -1101,6 +1102,10 @@ void LayoutRematerialization::rewriteSlice(SetVector<Value> &slice,
                                            DenseMap<Value, Attribute> &layout,
                                            ConvertLayoutOp convertOp,
                                            IRMapping &mapping) {
+  std::optional<bool> enableForLoopSupport =
+      mlir::triton::tools::isEnvValueBool(mlir::triton::tools::getStrEnv(
+          "TRITON_INTEL_REMOVELAYOUTCONVERSION_SUPPORT_FOR_LOOP"));
+
   SetVector<Operation *> opsToRewrite;
   // Keep track of yield operands that need to be duplicated.
   DenseMap<Operation *, SmallVector<int>> yieldOperandsMap;
@@ -1126,12 +1131,13 @@ void LayoutRematerialization::rewriteSlice(SetVector<Value> &slice,
         opsToRewrite.insert(ifOp.elseYield().getOperation());
         yieldOperandsMap[ifOp.elseYield()].push_back(operandIdx);
       }
-      if (auto forOp = v.getDefiningOp<scf::ForOp>()) {
-        unsigned operandIdx = cast<OpResult>(v).getResultNumber();
-        auto yieldOp = forOp.getBody()->getTerminator();
-        yieldOperandsMap[yieldOp].push_back(operandIdx);
-        opsToRewrite.insert(yieldOp);
-      }
+      if (enableForLoopSupport)
+        if (auto forOp = v.getDefiningOp<scf::ForOp>()) {
+          unsigned operandIdx = cast<OpResult>(v).getResultNumber();
+          auto yieldOp = forOp.getBody()->getTerminator();
+          yieldOperandsMap[yieldOp].push_back(operandIdx);
+          opsToRewrite.insert(yieldOp);
+        }
     } else {
       BlockArgument blockArg = cast<BlockArgument>(v);
       Operation *parentOp = blockArg.getOwner()->getParentOp();
@@ -1155,17 +1161,19 @@ void LayoutRematerialization::rewriteSlice(SetVector<Value> &slice,
   IRRewriter builder(slice.begin()->getContext());
   for (Operation *op : opsToRewrite) {
     if (auto forOp = dyn_cast<scf::ForOp>(op)) {
-      // Construct the new initialization argument by adding yielded operands
-      // that have been remapped.
-      auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
-      auto yieldOperands = llvm::to_vector(yieldOp.getOperands());
-      SmallVector<int> operandsToRewrite = yieldOperandsMap[yieldOp];
-      std::sort(operandsToRewrite.begin(), operandsToRewrite.end());
       SmallVector<Value> newOperands;
-      for (int operandIdx : operandsToRewrite) {
-        Value yieldOperand = yieldOp.getOperand(operandIdx);
-        if (mapping.contains(yieldOperand))
-          newOperands.push_back(mapping.lookup(yieldOperand));
+      if (enableForLoopSupport) {
+        // Construct the new initialization argument by adding yielded operands
+        // that have been remapped.
+        auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
+        auto yieldOperands = llvm::to_vector(yieldOp.getOperands());
+        SmallVector<int> operandsToRewrite = yieldOperandsMap[yieldOp];
+        std::sort(operandsToRewrite.begin(), operandsToRewrite.end());
+        for (int operandIdx : operandsToRewrite) {
+          Value yieldOperand = yieldOp.getOperand(operandIdx);
+          if (mapping.contains(yieldOperand))
+            newOperands.push_back(mapping.lookup(yieldOperand));
+        }
       }
 
       // Keep a mapping of the operands index to the new operands index.
@@ -1183,17 +1191,19 @@ void LayoutRematerialization::rewriteSlice(SetVector<Value> &slice,
       scf::ForOp newForOp = replaceForOpWithNewSignature(
           builder, forOp, newOperands, replacements);
 
-      // Add rematerializations for loop results in the slice.
-      unsigned oldIdx = 0;
-      unsigned newIdx = forOp.getNumResults();
-      for (auto res : forOp.getResults()) {
-        if (slice.count(res)) {
-          mapping.map(forOp.getResult(oldIdx), newForOp.getResult(newIdx));
-          addRematValue(forOp.getResult(oldIdx), layout[res],
-                        newForOp.getResult(newIdx));
-          ++newIdx;
+      if (enableForLoopSupport) {
+        // Add rematerializations for loop results in the slice.
+        unsigned oldIdx = 0;
+        unsigned newIdx = forOp.getNumResults();
+        for (auto res : forOp.getResults()) {
+          if (slice.count(res)) {
+            mapping.map(forOp.getResult(oldIdx), newForOp.getResult(newIdx));
+            addRematValue(forOp.getResult(oldIdx), layout[res],
+                          newForOp.getResult(newIdx));
+            ++newIdx;
+          }
+          ++oldIdx;
         }
-        ++oldIdx;
       }
 
       deadOps.push_back(forOp.getOperation());

third_party/intel/lib/TritonIntelGPUTransforms/Utility.cpp

@@ -15,6 +15,7 @@
 #include "intel/include/Dialect/TritonIntelGPU/Transforms/Utility.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
+#include "triton/Tools/Sys/GetEnv.hpp"
 
 #include <optional>
 
@@ -186,6 +187,10 @@ LogicalResult getConvertBackwardSlice(
   DenseSet<std::pair<OpOperand *, Attribute>> seen;
   SmallVector<std::pair<OpOperand *, Attribute>> queue;
 
+  std::optional<bool> enableForLoopSupport =
+      mlir::triton::tools::isEnvValueBool(mlir::triton::tools::getStrEnv(
+          "TRITON_INTEL_REMOVELAYOUTCONVERSION_SUPPORT_FOR_LOOP"));
+
   auto enqueue = [&](OpOperand &operand, Attribute encoding) {
     auto x = std::make_pair(&operand, encoding);
     if (!seen.insert(x).second) {
@@ -217,6 +222,11 @@ LogicalResult getConvertBackwardSlice(
     if (!isTensorOrTensorPointerType(currentValue.getType()))
       continue;
 
+    // Skip propagating through for op results for now.
+    // TODO: enable this based on needs.
+    if (!enableForLoopSupport && currentValue.getDefiningOp<scf::ForOp>())
+      return failure();
+
     if (failed(updateLayout(currentValue, encoding)))
       return failure();