Merge commit '815b2a46f633cb98459874c75bb74e79051a0f78'

Author: anmyachev

.github/workflows/integration-tests-amd.yml

@@ -31,7 +31,11 @@ jobs:
       CCACHE_COMPRESS: "true"
     container:
       image: ${{ matrix.image }}
-      options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
+      # Cache save/restore is on the host machine at directory /home/runner/.triton, while in the docker
+      # container expect it at /github/home/.triton. So map here to make sure visible in docker.
+      options: >-
+        --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
+        --volume /home/runner/.triton:/github/home/.triton
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -54,7 +58,6 @@ jobs:
           echo "llvm=$(cat $llvm_file | cut -c 1-8)" >> $GITHUB_OUTPUT
           echo "nvidia=$(sha256sum $nvidia_file | cut -d ' ' -f 1)" >> $GITHUB_OUTPUT
           echo "json=$(cat $json_file)" >> $GITHUB_OUTPUT
-          echo "datetime=$(date -u -Iseconds)" >> $GITHUB_OUTPUT
         shell: bash
       - name: Cache build dependencies
         uses: actions/cache@v4
@@ -162,5 +165,5 @@ jobs:
         # Always cleanup the worker, even if builds or tests failed
         if: always()
         run: |
-          rm -rf ~/.triton
+          rm -rf ~/.triton/cache
           rm -rf ~/.ccache

include/triton/Dialect/TritonGPU/Transforms/Passes.h

@@ -2,6 +2,7 @@
 #define TRITON_DIALECT_TRITONGPU_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
+#include "nvidia/include/Dialect/NVWS/IR/Dialect.h"
 #include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 namespace mlir {

include/triton/Dialect/TritonGPU/Transforms/Passes.td

@@ -106,7 +106,8 @@ def TritonGPUAutomaticWarpSpecialization : Pass<"tritongpu-automatic-warp-specia
     "mlir::triton::gpu::TritonGPUDialect",
     "mlir::scf::SCFDialect",
     "mlir::arith::ArithDialect",
-    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect"
+    "mlir::triton::nvidia_gpu::TritonNvidiaGPUDialect",
+    "triton::nvws::NVWSDialect"
   ];
 
   let options = [
@@ -143,7 +144,10 @@ def TritonGPUPartitionLoops : Pass<"tritongpu-partition-loops", "mlir::ModuleOp"
     between any of the partitions.
   }];
 
-  let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect"];
+  let dependentDialects = [
+    "mlir::triton::gpu::TritonGPUDialect",
+    "triton::nvws::NVWSDialect"
+  ];
 }
 
 def TritonGPUOptimizePartitionWarps : Pass<"tritongpu-optimize-partition-warps", "mlir::ModuleOp"> {

lib/Conversion/TritonGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -278,9 +278,8 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     assert(cvtNeedsSharedMemory(op.getSrc().getType(), op.getType()));
 
     // Try to use swizzling to implement the conversion
-    // HACK Remove once AMD tests pass for the swizzling path
-    if (targetInfo.isCuda() && succeeded(transferWithinBlockSwizzling(
-                                   op, adaptor.getSrc(), rewriter))) {
+    if (succeeded(
+            transferWithinBlockSwizzling(op, adaptor.getSrc(), rewriter))) {
       return success();
     }
 

lib/Dialect/Gluon/IR/Dialect.cpp

@@ -1,7 +1,7 @@
 #include "triton/Dialect/Gluon/IR/Dialect.h"
 
-#include "mlir/IR/DialectImplementation.h"
 #include "mlir/Support/LLVM.h"
+#include "triton/Dialect/Triton/IR/Interfaces.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
@@ -111,6 +111,7 @@ void GluonDialect::initialize() {
 #define GET_OP_LIST
 #include "triton/Dialect/Gluon/IR/Ops.cpp.inc"
       >();
+  addInterfaces<TritonInlinerInterface>();
   addInterfaces<GluonInferLayoutInterface>();
 }
 

lib/Dialect/Gluon/Transforms/ResolveAutoEncodings.cpp

@@ -66,12 +66,9 @@ LogicalResult inferAutoLayouts(FuncOp func) {
     return success();
   };
 
-  // 1. Set seed values from layout conversions
+  // 1. Set seed values from set_auto_layout ops
   auto res = func.walk([&](gluon::SetAutoLayoutOp op) -> WalkResult {
-    auto res = updateEncoding({op.getSrc()}, op.getType().getEncoding());
-    op.getResult().replaceAllUsesWith(op.getSrc());
-    op->erase();
-    return res;
+    return updateEncoding({op.getSrc()}, op.getType().getEncoding());
   });
 
   if (res.wasInterrupted())
@@ -153,6 +150,13 @@ LogicalResult inferAutoLayouts(FuncOp func) {
     }
   }
 
+  // 4. Cleanup set_auto_layout ops
+  func.walk([&](gluon::SetAutoLayoutOp op) {
+    assert(op.getSrc().getType() == op.getType());
+    op.getResult().replaceAllUsesWith(op.getSrc());
+    op->erase();
+  });
+
   return success();
 }
 

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -46,6 +46,8 @@ add_triton_library(TritonGPUTransforms
   TritonTransforms
   TritonGPUIR
   TritonNvidiaGPUIR
+  NVWSIR
+  NVWSTransforms
   TritonToTritonGPU
   TritonInstrumentIR
   MLIRTransformUtils