address comments and add test

triton-lang · Nov 6, 2024 · 70d72e8 · 70d72e8
1 parent f1b904c
commit 70d72e8
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/test/TritonGPU/amd/optimize-lds-usage.mlir b/test/TritonGPU/amd/optimize-lds-usage.mlir
@@ -119,6 +119,8 @@ module attributes {"triton_gpu.num-warps" = 8 : i32, "triton_gpu.threads-per-war
   }
 }
 
+// -----
+
 // Checks that optimization do not crash on 1d tensor
 // CHECK-LABEL: convert_1d
 // CHECK: triton_gpu.local_alloc

diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUsage.cpp
@@ -113,7 +113,7 @@ class OptimizeAMDLDSUsage
 
     // Special case for rank == 1
     if (rank == 1) {
-      threadsPerWarp[0] = warpSize;
+      threadsPerWarp[0] = warpSize / 8;
     } else {
       assert(rank > 1);
       threadsPerWarp[rank - 1] = warpSize / 8;

diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/OptimizeLDSUtility.cpp
@@ -69,7 +69,8 @@ Attribute createTmpLayout(Attribute layout, ArrayRef<unsigned> warpsPerCTA) {
         src.getKWidth());
   }
   if (auto src = dyn_cast<triton::gpu::SliceEncodingAttr>(layout)) {
-    // TODO: think of a way to construct slice layouts based on warpsPerCTA argument
+    // TODO: think of a way to construct slice layouts based on warpsPerCTA
+    // argument
     auto parentWarpsPerCTA = triton::gpu::getWarpsPerCTA(src.getParent());
     return triton::gpu::SliceEncodingAttr::get(
         ctx, src.getDim(), createTmpLayout(src.getParent(), parentWarpsPerCTA));