Fix conversion of loops to GPU with no block/thread dimensions.

Summary:
The current code assumes that one always maps at least one loop to block
dimensions and at least one loop to thread dimensions. If either is not
the case, a loop would get mapped twice.

Differential Revision: https://reviews.llvm.org/D73685
This commit is contained in:
Stephan Herhut 2020-01-30 11:27:17 +01:00
parent cf6037b561
commit 84695dd4d7
2 changed files with 36 additions and 2 deletions

View File

@ -403,10 +403,10 @@ void LoopToGpuConverter::createLaunch(OpTy rootForOp, OpTy innermostForOp,
Value constOne = (numBlockDims < 3 || numThreadDims < 3)
? builder.create<ConstantIndexOp>(rootForOp.getLoc(), 1)
: nullptr;
Value gridSizeX = dims[0];
Value gridSizeX = numBlockDims > 0 ? dims[0] : constOne;
Value gridSizeY = numBlockDims > 1 ? dims[1] : constOne;
Value gridSizeZ = numBlockDims > 2 ? dims[2] : constOne;
Value blockSizeX = dims[numBlockDims];
Value blockSizeX = numThreadDims > 0 ? dims[numBlockDims] : constOne;
Value blockSizeY = numThreadDims > 1 ? dims[numBlockDims + 1] : constOne;
Value blockSizeZ = numThreadDims > 2 ? dims[numBlockDims + 2] : constOne;

View File

@ -0,0 +1,34 @@
// RUN: mlir-opt -convert-loops-to-gpu -gpu-block-dims=0 -gpu-thread-dims=1 %s | FileCheck --check-prefix=CHECK-THREADS %s --dump-input-on-failure
// RUN: mlir-opt -convert-loops-to-gpu -gpu-block-dims=1 -gpu-thread-dims=0 %s | FileCheck --check-prefix=CHECK-BLOCKS %s --dump-input-on-failure
// CHECK-THREADS-LABEL: @one_d_loop
// CHECK-BLOCKS-LABEL: @one_d_loop
func @one_d_loop(%A : memref<?xf32>, %B : memref<?xf32>) {
// Bounds of the loop, its range and step.
// CHECK-THREADS-NEXT: %{{.*}} = constant 0 : index
// CHECK-THREADS-NEXT: %{{.*}} = constant 42 : index
// CHECK-THREADS-NEXT: %[[BOUND:.*]] = subi %{{.*}}, %{{.*}} : index
// CHECK-THREADS-NEXT: %{{.*}} = constant 1 : index
// CHECK-THREADS-NEXT: %[[ONE:.*]] = constant 1 : index
//
// CHECK-BLOCKS-NEXT: %{{.*}} = constant 0 : index
// CHECK-BLOCKS-NEXT: %{{.*}} = constant 42 : index
// CHECK-BLOCKS-NEXT: %[[BOUND:.*]] = subi %{{.*}}, %{{.*}} : index
// CHECK-BLOCKS-NEXT: %{{.*}} = constant 1 : index
// CHECK-BLOCKS-NEXT: %[[ONE:.*]] = constant 1 : index
// CHECK-THREADS-NEXT: gpu.launch blocks(%[[B0:.*]], %[[B1:.*]], %[[B2:.*]]) in (%{{.*}} = %[[ONE]], %{{.*}} = %[[ONE]], %{{.*}}0 = %[[ONE]]) threads(%[[T0:.*]], %[[T1:.*]], %[[T2:.*]]) in (%{{.*}} = %[[BOUND]], %{{.*}} = %[[ONE]], %{{.*}} = %[[ONE]])
// CHECK-BLOCKS-NEXT: gpu.launch blocks(%[[B0:.*]], %[[B1:.*]], %[[B2:.*]]) in (%{{.*}} = %[[BOUND]], %{{.*}} = %[[ONE]], %{{.*}}0 = %[[ONE]]) threads(%[[T0:.*]], %[[T1:.*]], %[[T2:.*]]) in (%{{.*}} = %[[ONE]], %{{.*}} = %[[ONE]], %{{.*}} = %[[ONE]])
affine.for %i = 0 to 42 {
// CHECK-THREADS-NEXT: %[[INDEX:.*]] = addi %{{.*}}, %[[T0]]
// CHECK-THREADS-NEXT: load %{{.*}}[%[[INDEX]]]
// CHECK-BLOCKS-NEXT: %[[INDEX:.*]] = addi %{{.*}}, %[[B0]]
// CHECK-BLOCKS-NEXT: load %{{.*}}[%[[INDEX]]]
%0 = load %A[%i] : memref<?xf32>
store %0, %B[%i] : memref<?xf32>
// CHECK-THREADS: gpu.terminator
// CHECK-BLOCKS: gpu.terminator
}
return
}