Fix allocation for paded shared layout

Author: zhanglx13

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1781,7 +1781,8 @@ int64_t PaddedSharedEncodingAttr::getPaddedSize(ArrayRef<int64_t> shape) const {
     paddingSize += (unpaddedSize >> llvm::Log2_32(interval))
                    << llvm::Log2_32(padding);
     // There is no need for padding after the last element
-    paddingSize -= padding;
+    if (unpaddedSize % interval == 0)
+      paddingSize -= padding;
   }
   return unpaddedSize + paddingSize;
 }

test/Analysis/test-allocation.mlir

@@ -947,26 +947,26 @@ tt.func @padded_shared_layout_size() {
   // expected-remark @+2 {{offset = 0, size = 510}}
   // 255 * 2B = 510B
   %alloc0 = ttg.local_alloc : () -> !ttg.memdesc<1x255xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 528}}
-  // (256 + 8) * 2B = 528B
+  // expected-remark @+2 {{offset = 0, size = 512}}
+  // 256 * 2B = 512B
   %alloc1 = ttg.local_alloc : () -> !ttg.memdesc<1x256xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
   // expected-remark @+2 {{offset = 0, size = 530}}
   // (257 + 8) * 2B = 530B
   %alloc2 = ttg.local_alloc : () -> !ttg.memdesc<1x257xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
   // expected-remark @+2 {{offset = 0, size = 1038}}
   // (511 + 8) * 2B = 1038B
   %alloc3 = ttg.local_alloc : () -> !ttg.memdesc<1x511xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 1056}}
-  // (512 + 8 * 2) * 2B = 1056B
+  // expected-remark @+2 {{offset = 0, size = 1040}}
+  // (512 + 8 * 1) * 2B = 1040B
   %alloc4 = ttg.local_alloc : () -> !ttg.memdesc<1x512xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
   // expected-remark @+2 {{offset = 0, size = 1058}}
   // (513 + 8 * 2) * 2B = 1058B
   %alloc5 = ttg.local_alloc : () -> !ttg.memdesc<1x513xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 528}}
-  // (16 * 16 + 8) * 2B = 528B
+  // expected-remark @+2 {{offset = 0, size = 512}}
+  // 16 * 16 * 2B = 512B
   %alloc6 = ttg.local_alloc : () -> !ttg.memdesc<16x16xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 1056}}
-  // (16 * 32 + 8 * 2) * 2B = 1056B
+  // expected-remark @+2 {{offset = 0, size = 1040}}
+  // (16 * 32 + 8 * 1) * 2B = 1040B
   %alloc7 = ttg.local_alloc : () -> !ttg.memdesc<16x32xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
   // expected-remark @+2 {{offset = 0, size = 1008}}
   // (31 * 16 + 8) * 2B = 1008B
@@ -975,28 +975,28 @@ tt.func @padded_shared_layout_size() {
 }
 
 // expected-remark @below {{padded_shared_layout_element_type}}
-// expected-remark @below {{size = 16896}}
+// expected-remark @below {{size = 16864}}
 tt.func @padded_shared_layout_element_type() {
-  // expected-remark @+2 {{offset = 0, size = 4224}}
-  // (16 * 256 + 8 * 16) * 1B = 4224B
+  // expected-remark @+2 {{offset = 0, size = 4216}}
+  // (16 * 256 + 8 * 15) * 1B = 4216B
   %alloc0 = ttg.local_alloc : () -> !ttg.memdesc<16x256xi8, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 8448}}
-  // (16 * 256 + 8 * 16) * 2B = 8448B
+  // expected-remark @+2 {{offset = 0, size = 8432}}
+  // (16 * 256 + 8 * 15) * 2B = 8432B
   %alloc1 = ttg.local_alloc : () -> !ttg.memdesc<16x256xf16, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 16896}}
-  // (16 * 256 + 8 * 16) * 4B = 16896B
+  // expected-remark @+2 {{offset = 0, size = 16864}}
+  // (16 * 256 + 8 * 15) * 4B = 16864B
   %alloc2 = ttg.local_alloc : () -> !ttg.memdesc<16x256xf32, #PADDED_SHARED_0, #ttg.shared_memory, mutable>
   tt.return
 }
 
 // expected-remark @below {{padded_shared_layout_multi_tier}}
-// expected-remark @below {{size = 4480}}
+// expected-remark @below {{size = 4466}}
 tt.func @padded_shared_layout_multi_tier() {
-  // expected-remark @+2 {{offset = 0, size = 4352}}
-  // (16 * 256 + 4 * 32 + 8 * 16) * 1B = 4352B
+  // expected-remark @+2 {{offset = 0, size = 4340}}
+  // (16 * 256 + 4 * 31 + 8 * 15) * 1B = 4340B
   %alloc0 = ttg.local_alloc : () -> !ttg.memdesc<16x256xi8, #PADDED_SHARED_1, #ttg.shared_memory, mutable>
-  // expected-remark @+2 {{offset = 0, size = 4480}}
-  // (16 * 256 + 2 * 64 + 4 * 32 + 8 * 16) * 1B = 4480B
+  // expected-remark @+2 {{offset = 0, size = 4466}}
+  // (16 * 256 + 2 * 63 + 4 * 31 + 8 * 15) * 1B = 4466B
   %alloc1 = ttg.local_alloc : () -> !ttg.memdesc<16x256xi8, #PADDED_SHARED_2, #ttg.shared_memory, mutable>
   tt.return
 }