[CUDA] Add some basic logic to pick better configuration (#6754) Pick configuration based on the problem size and skip using shared memory for workgroup of a single warp.
diff --git a/iree/compiler/Codegen/LLVMGPU/BUILD b/iree/compiler/Codegen/LLVMGPU/BUILD index ae2484a..0012dc6 100644 --- a/iree/compiler/Codegen/LLVMGPU/BUILD +++ b/iree/compiler/Codegen/LLVMGPU/BUILD
@@ -36,6 +36,7 @@ "//iree/compiler/Codegen/Common", "//iree/compiler/Codegen/Transforms", "//iree/compiler/Codegen/Utils", + "//iree/compiler/Dialect/Flow/IR", "//iree/compiler/Dialect/HAL/IR", "//iree/compiler/Dialect/LinalgExt/IR", "//iree/compiler/Dialect/LinalgExt/Transforms",
diff --git a/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt b/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt index 3b95579..9540af2 100644 --- a/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt +++ b/iree/compiler/Codegen/LLVMGPU/CMakeLists.txt
@@ -63,6 +63,7 @@ iree::compiler::Codegen::PassHeaders iree::compiler::Codegen::Transforms iree::compiler::Codegen::Utils + iree::compiler::Dialect::Flow::IR iree::compiler::Dialect::HAL::IR iree::compiler::Dialect::LinalgExt::IR iree::compiler::Dialect::LinalgExt::Transforms
diff --git a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp index c41fcb9..a3d2fef 100644 --- a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp +++ b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -7,7 +7,9 @@ #include "iree/compiler/Codegen/LLVMGPU/KernelConfig.h" #include "iree/compiler/Codegen/Utils/Utils.h" +#include "iree/compiler/Dialect/Flow/IR/FlowOps.h" #include "llvm/Support/Debug.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" #include "mlir/IR/Types.h" #include "mlir/IR/Value.h" @@ -30,8 +32,10 @@ tileSizes.push_back(TileWorkgroupSizePair({{64, 128, 8}, {32, 4, 1}})); tileSizes.push_back(TileWorkgroupSizePair({{128, 64, 8}, {16, 8, 1}})); tileSizes.push_back(TileWorkgroupSizePair({{16, 256, 16}, {64, 2, 1}})); + tileSizes.push_back(TileWorkgroupSizePair({{8, 128, 4}, {32, 1, 1}})); tileSizes.push_back(TileWorkgroupSizePair({{16, 64, 4}, {16, 2, 1}})); + tileSizes.push_back(TileWorkgroupSizePair({{1, 128, 8}, {32, 1, 1}})); } static LogicalResult setContractConfig(FuncOp entryPoint, linalg::LinalgOp op) { @@ -61,6 +65,12 @@ int64_t tileY = 256; int64_t tileK = 4; SmallVector<int64_t, 3> workgroupSize = {2 * cudaWarpSize, 1, 1}; + // Special case for very small matrices. + if (sizeM * sizeN <= cudaWarpSize) { + tileX = sizeN; + tileY = sizeM; + workgroupSize = {sizeM, sizeN, 1}; + } SmallVector<TileWorkgroupSizePair> tileSizeConfig; // Query the best configuration. getMatmulConfig(tileSizeConfig); @@ -114,7 +124,7 @@ size_t numLoops = partitionedLoops.back() + 1; std::array<int64_t, 3> workgroupSize = {cudaWarpSize, 1, 1}; - static constexpr unsigned vectorSize = 4; + unsigned vectorSize = 4; SmallVector<int64_t, 4> workgroupTileSizes(numLoops, 1), threadTileSizes(numLoops, 1); // Set all non-parallel loops to zero tile size. @@ -126,6 +136,36 @@ threadTileSizes[depth] = 0; } } + + auto genericOp = dyn_cast<linalg::GenericOp>(op); + bool outputSizeIsProblemSize = + genericOp && + llvm::all_of(genericOp.getOutputOperands(), + [&genericOp](OpOperand *outputOperand) { + return genericOp.getTiedIndexingMap(outputOperand) + .isProjectedPermutation(); + }); + if (outputSizeIsProblemSize) { + // Calculate the problem size to adjust the tile size. + int64_t problemSize = 1; + entryPoint.walk([&problemSize](IREE::Flow::DispatchTensorStoreOp storeOp) { + ArrayRef<int64_t> shape = storeOp.target() + .getType() + .cast<IREE::Flow::DispatchTensorType>() + .getShape(); + int64_t prod = 1; + for (int64_t dim : shape) prod *= dim; + problemSize = std::max(prod, problemSize); + }); + // If the problem size is too small or if the op cannot be vectorized, + // reduce the vector size to prevent bad memory access patterns. + if ((problemSize / (cudaWarpSize * vectorSize)) < 64) vectorSize = 1; + } + // Pick a vectorSize of 1 for op that we know won't get vectorizedd. + // TODO(thomasraoux): This could be improved by checking if the linalg op + // would fail vectorization. + if (!isa<linalg::LinalgOp>(op)) vectorSize = 1; + // Set the inner most parallel loop to `lowerTs`. for (int64_t depth = numLoops; depth > 0; depth--) { if (partitionedLoopsSet.count(depth - 1)) {
diff --git a/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp b/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp index 18a69b2..4d66926 100644 --- a/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp +++ b/iree/compiler/Codegen/LLVMGPU/LLVMGPUTileAndDistribute.cpp
@@ -181,7 +181,8 @@ static void populatePromotionPatterns(MLIRContext *context, OwningRewritePatternList &patterns) { - patterns.insert<linalg::LinalgPromotionPattern<linalg::MatmulOp>>( + patterns.insert<linalg::LinalgPromotionPattern<linalg::MatmulOp>, + linalg::LinalgPromotionPattern<linalg::BatchMatmulOp>>( context, linalg::LinalgPromotionOptions() .setAllocationDeallocationFns(allocateWorkgroupMemory, @@ -228,7 +229,13 @@ funcOp.dump(); }); - { + auto workgroupSize = llvm::to_vector<4>(llvm::map_range( + getEntryPoint(funcOp).workgroup_size().getValue(), + [&](Attribute attr) { return attr.cast<IntegerAttr>().getInt(); })); + int64_t flatWorkgroupSize = + workgroupSize[0] * workgroupSize[1] * workgroupSize[2]; + // Only promote to workgroup size if there are multiple warps. + if (flatWorkgroupSize > 32) { OwningRewritePatternList promotionPatterns(&getContext()); populatePromotionPatterns(context, promotionPatterns); (void)applyPatternsAndFoldGreedily(funcOp, std::move(promotionPatterns)); @@ -264,9 +271,6 @@ }); { - auto workgroupSize = llvm::to_vector<4>(llvm::map_range( - getEntryPoint(funcOp).workgroup_size().getValue(), - [&](Attribute attr) { return attr.cast<IntegerAttr>().getInt(); })); // Apply last level of tiling and distribute to threads. OwningRewritePatternList threadLevelTilingPatterns(context); populateTilingToInvocationPatterns(context, threadLevelTilingPatterns,
diff --git a/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir b/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir index 4f8c6c6..c5d0979 100644 --- a/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir +++ b/iree/compiler/Codegen/LLVMGPU/test/gpu_set_num_workgroups.mlir
@@ -1,50 +1,40 @@ // RUN: iree-opt -split-input-file -pass-pipeline='hal.executable(hal.executable.variant(iree-llvmgpu-lower-executable-target-pass{test-lowering-configuration}))' %s | IreeFileCheck %s -#executable_target_cuda_nvptx_fb = #hal.executable.target<"cuda", "cuda-nvptx-fb"> -#map0 = affine_map<()[s0, s1] -> (s0 * s1)> -#map1 = affine_map<(d0)[s0] -> (d0 + s0)> -#map2 = affine_map<(d0) -> (d0)> -#map3 = affine_map<(d0)[s0] -> (s0, -d0 + 1024)> -hal.executable @add_dispatch_0 attributes {sym_visibility = "private"} { +hal.executable @add_dispatch_0 { + hal.interface @io { + hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @ret0, set=0, binding=1, type="StorageBuffer", access="Write|Discard" + } hal.executable.variant @cuda, target = #hal.executable.target<"cuda", "cuda-nvptx-fb"> { - hal.executable.entry_point @add_dispatch_0 attributes {interface = @io, ordinal = 0 : index} - module { - func @add_dispatch_0() { - %c0 = constant 0 : index - %c1024 = constant 1024 : index - %0 = hal.interface.binding.subspan @io::@ro0[%c0] : memref<1024xf32> - %1 = hal.interface.binding.subspan @io::@ro1[%c0] : memref<1024xf32> - %2 = hal.interface.binding.subspan @io::@wo2[%c0] : memref<1024xf32> - %workgroup_size_x = hal.interface.workgroup.size[0] : index - %workgroup_id_x = hal.interface.workgroup.id[0] : index - %workgroup_count_x = hal.interface.workgroup.count[0] : index - %3 = affine.apply #map0()[%workgroup_id_x, %workgroup_size_x] - %4 = affine.apply #map0()[%workgroup_count_x, %workgroup_size_x] - scf.for %arg0 = %3 to %c1024 step %4 { - %5 = affine.min #map3(%arg0)[%workgroup_size_x] - %6 = memref.subview %0[%arg0] [%5] [1] : memref<1024xf32> to memref<?xf32, #map1> - %7 = memref.subview %1[%arg0] [%5] [1] : memref<1024xf32> to memref<?xf32, #map1> - %8 = memref.subview %2[%arg0] [%5] [1] : memref<1024xf32> to memref<?xf32, #map1> - linalg.generic { - indexing_maps = [#map2, #map2, #map2], - iterator_types = ["parallel"]} - ins(%6, %7 : memref<?xf32, #map1>, memref<?xf32, #map1>) - outs(%8 : memref<?xf32, #map1>) attrs = {__internal_linalg_transform__ = "workgroup"} { - ^bb0(%arg1: f32, %arg2: f32, %arg3: f32): // no predecessors - %9 = addf %arg1, %arg2 : f32 - linalg.yield %9 : f32 - } - } + hal.executable.entry_point @add_dispatch_0 attributes {interface = @io, ordinal = 0 : index} + module { + func @add_dispatch_0() { + %c0 = constant 0 : index + %0 = hal.interface.binding.subspan @io::@arg0[%c0] : !flow.dispatch.tensor<readonly:16384xf32> + %1 = hal.interface.binding.subspan @io::@arg1[%c0] : !flow.dispatch.tensor<readonly:16384xf32> + %2 = hal.interface.binding.subspan @io::@ret0[%c0] : !flow.dispatch.tensor<writeonly:16384xf32> + %3 = linalg.init_tensor [16384] : tensor<16384xf32> + %4 = flow.dispatch.tensor.load %0, offsets=[], sizes=[], strides=[] : !flow.dispatch.tensor<readonly:16384xf32> -> tensor<16384xf32> + %5 = flow.dispatch.tensor.load %1, offsets=[], sizes=[], strides=[] : !flow.dispatch.tensor<readonly:16384xf32> -> tensor<16384xf32> + %6 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%4, %5 : tensor<16384xf32>, tensor<16384xf32>) outs(%3 : tensor<16384xf32>) { + ^bb0(%arg0: f32, %arg1: f32, %arg2: f32): // no predecessors + %7 = addf %arg0, %arg1 : f32 + linalg.yield %7 : f32 + } -> tensor<16384xf32> + flow.dispatch.tensor.store %6, %2, offsets=[], sizes=[], strides=[] : tensor<16384xf32> -> !flow.dispatch.tensor<writeonly:16384xf32> return } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @arg0, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @arg1, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @ret0, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } } } } + // CHECK-DAG: #[[CONFIG:.+]] = {tileSizes = {{\[}}[128], [], [4]{{\]}}} // CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 128)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 * 128)> -// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0)[s0] -> (d0 + s0)> -// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0) -> (d0)> // CHECK: hal.executable.entry_point @add_dispatch_0 // CHECK-SAME: passPipeline = 3 : i32 // CHECK-SAME: workloadPerWorkgroup = [128] @@ -102,13 +92,13 @@ } } } -// CHECK-DAG: #[[CONFIG:.+]] = {tileSizes = {{\[}}[2, 256, 4], [], [2, 4]{{\]}}} -// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 256)> -// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)> +// CHECK-DAG: #[[CONFIG:.+]] = {tileSizes = {{\[}}[4, 2, 4], [], [1, 1]{{\]}}} +// CHECK-DAG: #[[MAP0:.+]] = affine_map<()[s0] -> (s0 ceildiv 2)> +// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0] -> (s0 ceildiv 4)> // CHECK: hal.executable.entry_point @dot_dispatch_1 // CHECK-SAME: passPipeline = 4 : i32 -// CHECK-SAME: workloadPerWorkgroup = [256, 2] -// CHECK-SAME: workgroup_size = [64 : index, 1 : index, 1 : index] +// CHECK-SAME: workloadPerWorkgroup = [2, 4] +// CHECK-SAME: workgroup_size = [2 : index, 4 : index, 1 : index] // CHECK-NEXT: ^bb0(%[[ARG0:[a-zA-Z0-9]+]]: index, %[[ARG1:[a-zA-Z0-9]+]]: index, // CHECK-DAG: %[[C1:.+]] = constant 1 : index // CHECK-DAG: %[[NWGS_X:.+]] = affine.apply #[[MAP0]]()[%[[ARG0]]]
diff --git a/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir b/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir index 8531b12..5947142 100644 --- a/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir +++ b/iree/compiler/Codegen/LLVMGPU/test/nvvm_pipeline_test.mlir
@@ -364,28 +364,28 @@ builtin.module { builtin.func @vector_add_dispatch() { %c0 = constant 0 : index - %c128 = constant 128 : index - %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:128xf32> - %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:128xf32> - %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:128xf32> + %c16384 = constant 16384 : index + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:16384xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:16384xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:16384xf32> %workgroup_size_x = hal.interface.workgroup.size[0] : index %workgroup_id_x = hal.interface.workgroup.id[0] : index %workgroup_count_x = hal.interface.workgroup.count[0] : index %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] - scf.for %arg0 = %3 to %c128 step %4 { - %5 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 128)>(%arg0)[%workgroup_size_x] - %6 = flow.dispatch.tensor.load %0, offsets = [%arg0], sizes = [%5], strides = [1] : !flow.dispatch.tensor<readonly:128xf32> -> tensor<?xf32> - %7 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 128)>(%arg0)[%workgroup_size_x] - %8 = flow.dispatch.tensor.load %1, offsets = [%arg0], sizes = [%7], strides = [1] : !flow.dispatch.tensor<readonly:128xf32> -> tensor<?xf32> - %9 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 128)>(%arg0)[%workgroup_size_x] + scf.for %arg0 = %3 to %c16384 step %4 { + %5 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 16384)>(%arg0)[%workgroup_size_x] + %6 = flow.dispatch.tensor.load %0, offsets = [%arg0], sizes = [%5], strides = [1] : !flow.dispatch.tensor<readonly:16384xf32> -> tensor<?xf32> + %7 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 16384)>(%arg0)[%workgroup_size_x] + %8 = flow.dispatch.tensor.load %1, offsets = [%arg0], sizes = [%7], strides = [1] : !flow.dispatch.tensor<readonly:16384xf32> -> tensor<?xf32> + %9 = affine.min affine_map<(d0, d1) -> (d1, -d0 + 16384)>(%arg0)[%workgroup_size_x] %10 = linalg.init_tensor [%9] : tensor<?xf32> %11 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%6, %8 : tensor<?xf32>, tensor<?xf32>) outs(%10 : tensor<?xf32>) attrs = {__internal_linalg_transform__ = "workgroup"} { ^bb0(%arg1: f32, %arg2: f32, %arg3: f32): // no predecessors %12 = addf %arg1, %arg2 : f32 linalg.yield %12 : f32 } -> tensor<?xf32> - flow.dispatch.tensor.store %11, %2, offsets = [%arg0], sizes = [%9], strides = [1] : tensor<?xf32> -> !flow.dispatch.tensor<writeonly:128xf32> + flow.dispatch.tensor.store %11, %2, offsets = [%arg0], sizes = [%9], strides = [1] : tensor<?xf32> -> !flow.dispatch.tensor<writeonly:16384xf32> } return } @@ -406,8 +406,8 @@ // ----- #map0 = affine_map<()[s0, s1] -> (s0 * s1)> -#map1 = affine_map<(d0)[s0] -> (s0, -d0 + 768)> -#map2 = affine_map<(d0)[s0] -> (-d0 + 768, s0)> +#map1 = affine_map<(d0)[s0] -> (s0, -d0 + 16384)> +#map2 = affine_map<(d0)[s0] -> (-d0 + 16384, s0)> #map3 = affine_map<(d0, d1) -> (d1, d0)> #map4 = affine_map<(d0, d1) -> (d0)> @@ -417,18 +417,18 @@ builtin.module { builtin.func @vector_reduction_dispatch() { %c0 = constant 0 : index - %c768 = constant 768 : index + %c16384 = constant 16384 : index %cst = constant 1.000000e+00 : f32 - %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:512x768xf32> - %1 = hal.interface.binding.subspan @io::@s0b1_xw_external[%c0] : !flow.dispatch.tensor<writeonly:768xf32> + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:512x16384xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_xw_external[%c0] : !flow.dispatch.tensor<writeonly:16384xf32> %workgroup_size_x = hal.interface.workgroup.size[0] : index %workgroup_id_x = hal.interface.workgroup.id[0] : index %workgroup_count_x = hal.interface.workgroup.count[0] : index %2 = affine.apply #map0()[%workgroup_id_x, %workgroup_size_x] %3 = affine.apply #map0()[%workgroup_count_x, %workgroup_size_x] - scf.for %arg0 = %2 to %c768 step %3 { + scf.for %arg0 = %2 to %c16384 step %3 { %4 = affine.min #map1(%arg0)[%workgroup_size_x] - %5 = flow.dispatch.tensor.load %0, offsets = [0, %arg0], sizes = [512, %4], strides = [1, 1] : !flow.dispatch.tensor<readonly:512x768xf32> -> tensor<512x?xf32> + %5 = flow.dispatch.tensor.load %0, offsets = [0, %arg0], sizes = [512, %4], strides = [1, 1] : !flow.dispatch.tensor<readonly:512x16384xf32> -> tensor<512x?xf32> %6 = affine.min #map1(%arg0)[%workgroup_size_x] %7 = affine.min #map2(%arg0)[%workgroup_size_x] %8 = linalg.init_tensor [%7] : tensor<?xf32> @@ -438,7 +438,7 @@ %11 = addf %arg1, %arg2 : f32 linalg.yield %11 : f32 } -> tensor<?xf32> - flow.dispatch.tensor.store %10, %1, offsets = [%arg0], sizes = [%6], strides = [1] : tensor<?xf32> -> !flow.dispatch.tensor<writeonly:768xf32> + flow.dispatch.tensor.store %10, %1, offsets = [%arg0], sizes = [%6], strides = [1] : tensor<?xf32> -> !flow.dispatch.tensor<writeonly:16384xf32> } return }