|
1 | | -diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp |
2 | | ---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp |
3 | | -+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp |
4 | | -@@ -1088,6 +1088,11 @@ struct AsyncCopyGlobalToLocalOpConversio |
5 | | - // mask bits are the same. For example if N=2, the mask must be |
6 | | - // [x, x, y, y, ...]. |
7 | | - unsigned maxVec = getContiguity(op.getSrc()); |
8 | | -+ |
9 | | -+ if (resElemTy.isF16() && maxVec == 1) { |
10 | | -+ maxVec = 2; |
11 | | -+ } |
12 | | -+ |
13 | | - if (mask) { |
14 | | - maxVec = std::min(maxVec, getMaskAlignment(mask)); |
15 | | - } |
| 1 | +diff --git a/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp b/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp |
| 2 | +--- a/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp |
| 3 | ++++ b/lib/Dialect/TritonGPU/Transforms/Pipeliner/LowerLoops.cpp |
| 4 | +@@ -508,7 +508,11 @@ scf::ForOp lowerLoads(scf::ForOp forOp, |
| 5 | + // bytes) |
| 6 | + int copyVecBytes = getCopyVecBytes( |
| 7 | + cast<RankedTensorType>(op.getResultTypes()[0]), sharedEncoding); |
| 8 | +- if (copyVecBytes >= 4 || isTMALoad(&op)) { |
| 9 | ++ bool is16BitElementOnBlackwell = |
| 10 | ++ (cast<RankedTensorType>(op.getResultTypes()[0]) |
| 11 | ++ .getElementTypeBitWidth() == 16) && |
| 12 | ++ getNVIDIAComputeCapability(forOp->getParentOfType<ModuleOp>()) >= 100; |
| 13 | ++ if ((copyVecBytes >= 4 && !is16BitElementOnBlackwell) || isTMALoad(&op)) { |
| 14 | + if (loadRequiresAdditionalBuffer(&op)) { |
| 15 | + // Allocate additional buffer required by the wgmma pipelining. |
| 16 | + stageDiff += 1; |
| 17 | +diff --git a/test/TritonGPU/loop-pipeline-hopper.mlir b/test/TritonGPU/loop-pipeline-hopper.mlir |
| 18 | +--- a/test/TritonGPU/loop-pipeline-hopper.mlir |
| 19 | ++++ b/test/TritonGPU/loop-pipeline-hopper.mlir |
| 20 | +@@ -52,7 +52,7 @@ |
| 21 | + // CHECK: %[[BSUB3:.*]] = ttg.memdesc_subview %[[BBUFFER]][%[[INS_IDX_3]], %[[CONSTANT_0]], %[[CONSTANT_0]]] |
| 22 | + // CHECK: %[[NEXT_B_BUFFER:.*]] = ttg.async_copy_global_to_local {{.*}}, %[[BSUB3]] |
| 23 | + // CHECK: scf.yield {{.*}}, %[[INS_IDX_3]], %[[EXT_IDX_3]] |
| 24 | +-module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { |
| 25 | ++module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90"} { |
| 26 | + tt.func @matmul_loop(%lb : index, %ub : index, %step : index, |
| 27 | + %A : !tt.ptr<f16> {tt.divisibility = 16 : i32}, |
| 28 | + %B : !tt.ptr<f16> {tt.divisibility = 16 : i32}) { |
| 29 | +@@ -587,7 +587,7 @@ module attributes {"ttg.num-ctas" = 1 : |
| 30 | + #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> |
| 31 | + #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> |
| 32 | + #smem = #ttg.shared_memory |
| 33 | +-module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { |
| 34 | ++module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90"} { |
| 35 | + // COMMON-LABEL: dot_prologue_epilogue |
| 36 | + // COMMON: {{.*}}, {{.*}}, %[[EXT:.*]]: i32, {{.*}} |
| 37 | + tt.func @dot_prologue_epilogue(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %ext: i32, %inc: tensor<64x16xi32, #blocked> {tt.divisibility = 16 : i32}) -> tensor<128x16xf32, #mma1> { |
| 38 | +@@ -660,7 +660,7 @@ module attributes {"ttg.num-ctas" = 1 : |
| 39 | + #shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}> |
| 40 | + #shared1 = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = true, elementBitWidth = 16}> |
| 41 | + #smem = #ttg.shared_memory |
| 42 | +-module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} { |
| 43 | ++module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.target" = "cuda:90"} { |
| 44 | + // CHECK-NOCANON-LABEL: pipeline_downstream_dependencies |
| 45 | + // CHECK-NOCANON: {{.*}}, {{.*}}, %[[EXT:.*]]: i32, {{.*}} |
| 46 | + tt.func @pipeline_downstream_dependencies(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %ext: i32, %inc: tensor<64x16xi32, #blocked> {tt.divisibility = 16 : i32}) -> tensor<128x16xf32, #mma1> { |
0 commit comments