[ROCm] Fixed an issue with InstructionSchedHintsPass

ROCm · Dec 15, 2024 · 2bd6f7e · 2bd6f7e
1 parent 2d9229a
commit 2bd6f7e
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 1 deletion.
diff --git a/third_party/triton/temporary/fix_InsertInstructionSchedHints.patch b/third_party/triton/temporary/fix_InsertInstructionSchedHints.patch
@@ -0,0 +1,11 @@
+--- a/third_party/amd/include/TritonAMDGPUToLLVM/Passes.td
++++ b/third_party/amd/include/TritonAMDGPUToLLVM/Passes.td
+@@ -59,7 +59,7 @@
+     let summary = "Insert instruction scheduling hints after the dot ops in the main loop";
+     let constructor = "mlir::triton::createInsertInstructionSchedHintsPass()";
+
+-    let dependentDialects = ["mlir::LLVM::LLVMDialect"];
++    let dependentDialects = ["mlir::LLVM::LLVMDialect", "mlir::triton::amdgpu::TritonAMDGPUDialect"];
+ }
+
+ def LowerInstructionSchedHints : Pass<"lower-insert-instruction-sched-hints", "mlir::ModuleOp"> {
diff --git a/third_party/triton/temporary/series.bzl b/third_party/triton/temporary/series.bzl
@@ -17,5 +17,6 @@ temporary_patch_list = [
     "//third_party/triton:temporary/fix_left_shift_overflow.patch",
     "//third_party/triton:temporary/prefetch.patch",
     "//third_party/triton:temporary/i4_to_bf16.patch",
+    "//third_party/triton:temporary/fix_InsertInstructionSchedHints.patch",
     # Add new patches just above this line
 ]
diff --git a/xla/service/gpu/fusions/triton/compilation_pipeline_rocm.cc b/xla/service/gpu/fusions/triton/compilation_pipeline_rocm.cc
@@ -91,7 +91,7 @@ absl::Status CreateTritonPipeline(
     pm.addPass(mlir::createTritonAMDGPUStreamPipelinePass());
     pm.addPass(mlir::createCanonicalizerPass());
   }
-  //pm.addPass(mt::createInsertInstructionSchedHintsPass());
+  pm.addPass(mt::createInsertInstructionSchedHintsPass());
   pm.addPass(mt::gpu::createTritonGPUOptimizeDotOperands({true}));
   pm.addPass(mt::gpu::createTritonGPURemoveLayoutConversions());
   pm.addPass(mt::gpu::createTritonGPUReduceDataDuplication());