[AMD] Added TRITON_HIP_STREAM_PREFETCH to GetEnv.hpp: invalidate cache

triton-lang · Nov 14, 2024 · bb70e36 · bb70e36
1 parent 22856a7
commit bb70e36
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/include/triton/Tools/Sys/GetEnv.hpp b/include/triton/Tools/Sys/GetEnv.hpp
@@ -28,6 +28,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "TRITON_DISABLE_LINE_INFO",
     "TRITON_DISABLE_RESHAPE_ENCODING_INFERENCE",
     "TRITON_ENABLE_LLVM_DEBUG",
+    "TRITON_HIP_STREAM_PREFETCH",
     "TRITON_LLVM_DEBUG_ONLY",
     "USE_IR_LOC",
     "NVPTX_ENABLE_DUMP",

diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/SchedInstructions.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/SchedInstructions.cpp
@@ -250,9 +250,9 @@ struct InstructionSchedHintsRewriter
   // local (LDS to registers) and global (HBM to registers) data prefetching.
   // see:
   // include/ck/tensor_operation/gpu/block/blockwise_gemm_pipeline_xdlops_v3.h
-  void
-  createCKV3Schedule(PatternRewriter &rewriter, Location loc,
-                     triton::amdgpu::InstructionSchedHint schedHint) const {
+  void createLocalPrefetchSchedule(
+      PatternRewriter &rewriter, Location loc,
+      triton::amdgpu::InstructionSchedHint schedHint) const {
 
     if (!(schedHint.getIsBufferLoadsAEnabled() &&
           schedHint.getIsBufferLoadsBEnabled())) {
@@ -456,7 +456,7 @@ struct InstructionSchedHintsRewriter
       break;
     }
     case SchedulingType::LOCAL_PREFETCH: {
-      createCKV3Schedule(rewriter, loc, instructionSchedHint);
+      createLocalPrefetchSchedule(rewriter, loc, instructionSchedHint);
       break;
     }
     case SchedulingType::NONE: