Fix bug of getShapePerCTATile for slice layout

Author: chengjunlu

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -209,12 +209,15 @@ SmallVector<unsigned> getShapePerCTATile(Attribute layout) {
           mlir::dyn_cast<DistributedEncodingTrait>(layout)) {
     auto sizePerThread = distributedLayout.getSizePerThread();
     auto threadsPerWarp = distributedLayout.getThreadsPerWarp();
-    // ThreadsPerWarp does not align with this function for slice layout
+    auto warpsPerCTA = distributedLayout.getWarpsPerCTA();
+    // ThreadsPerWarp and warpsPerCTA does not align with this function for
+    // slice layout
     if (auto sliceLayout = mlir::dyn_cast<SliceEncodingAttr>(layout)) {
       threadsPerWarp = getThreadsPerWarp(sliceLayout.getParent());
       threadsPerWarp.erase(threadsPerWarp.begin() + sliceLayout.getDim());
+      warpsPerCTA = getWarpsPerCTA(sliceLayout.getParent());
+      warpsPerCTA.erase(warpsPerCTA.begin() + sliceLayout.getDim());
     }
-    auto warpsPerCTA = distributedLayout.getWarpsPerCTA();
     assert(sizePerThread.size() == threadsPerWarp.size() &&
            sizePerThread.size() == warpsPerCTA.size());
     SmallVector<unsigned> shape;