Improve the heuristic on picking distributed size. (#8982) Fix #8980
diff --git a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp index 86aef61..4bed428 100644 --- a/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp +++ b/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -160,7 +160,8 @@ /// Flow level. static SmallVector<int64_t> getDefaultDistributedLoopTileSizes( ArrayRef<int64_t> lbs, ArrayRef<int64_t> ubs, - ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes) { + ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes, + ArrayRef<int64_t> vectorSizeHints) { assert(lbs.size() == ubs.size() && lbs.size() == minTileSizes.size() && lbs.size() == maxTileSizes.size() && "expected all vectors to be of equal size"); @@ -181,10 +182,23 @@ } int64_t candidateTileSize = 1; if (ubs[i] > lbs[i]) { - // Pick a value that evenly distributes the workload. - candidateTileSize = std::max<int64_t>( - llvm::PowerOf2Floor(static_cast<uint64_t>(ubs[i] - lbs[i]) / 2), - minTileSizes[i]); + int64_t dimSize = static_cast<uint64_t>(ubs[i] - lbs[i]); + int64_t targetSize = std::min(dimSize / 2, maxTileSizes[i]); + int64_t vectorSize = vectorSizeHints[i]; + if (vectorSize > 1) { + // Pick the factor of dim which is closest to the target tile size and + // is a multiplier of vector size. + for (int64_t k = vectorSize; k <= targetSize; k += vectorSize) { + if (dimSize % k == 0 && k >= minTileSizes[i]) { + candidateTileSize = k; + } + } + } + // Fallback to power of 2 if there's no hint or can't find the ideal size. + if (vectorSize <= 1 || candidateTileSize == 1) { + candidateTileSize = + std::max<int64_t>(llvm::PowerOf2Floor(targetSize), minTileSizes[i]); + } } // Limit the workload per workgroup to the default being the max to keep the @@ -205,18 +219,29 @@ } unsigned currDim = numDims; while (numWorkgroups > numWorkgroupsLimit && currDim > 0) { - if (distributedTileSizes[currDim - 1] >= maxTileSizes[currDim - 1] || - workload[currDim - 1] == ShapedType::kDynamicSize || - distributedTileSizes[currDim - 1] >= workload[currDim - 1]) { + unsigned index = currDim - 1; + int64_t currSize = distributedTileSizes[index]; + if (currSize >= maxTileSizes[index] || + workload[index] == ShapedType::kDynamicSize || + currSize >= workload[index]) { currDim--; continue; } - distributedTileSizes[currDim - 1] = std::min<int64_t>( - distributedTileSizes[currDim - 1] * 2, maxTileSizes[currDim - 1]); - int64_t nwg = - ceilFn(workload[currDim - 1], distributedTileSizes[currDim - 1]); - if (nwg < numWorkgroupsPerDim[currDim - 1]) { - numWorkgroups /= numWorkgroupsPerDim[currDim - 1]; + int64_t newSize = std::min<int64_t>(currSize * 2, maxTileSizes[index]); + int64_t vectorSize = vectorSizeHints[index]; + // Chech if it's the ideal size with vector size hint. And skip if the new + // size will break the ideal size. + if (vectorSize > 1 && + (currSize % vectorSize == 0 && workload[index] % currSize == 0) && + (newSize % vectorSize != 0 || workload[index] % newSize != 0)) { + currDim--; + continue; + } + + distributedTileSizes[index] = newSize; + int64_t nwg = ceilFn(workload[index], distributedTileSizes[index]); + if (nwg < numWorkgroupsPerDim[index]) { + numWorkgroups /= numWorkgroupsPerDim[index]; numWorkgroups *= nwg; } else { currDim--; @@ -253,12 +278,22 @@ /// Returns the tile size to use for the Flow level of an operation that /// implements the `PartitionableLoopsInterface`. +/// +/// The vectorSizeHints can be empty or as many as the number of loops. When not +/// empty, each hint should be 1 or the vector size. On the dimensions where the +/// hints != 1, it will try to find the tile sizes which are multipliers of the +/// hints. static SmallVector<int64_t> getDefaultDistributedLevelTileSizes( ArrayRef<Range> iterationDomain, IREE::Flow::PartitionableLoopsInterface partitionableLoopInterfaceOp, - ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes) { + ArrayRef<int64_t> minTileSizes, ArrayRef<int64_t> maxTileSizes, + ArrayRef<int64_t> vectorSizeHints = {}) { assert(iterationDomain.size() == minTileSizes.size() && "expected as many min tile sizes as number of loops"); + assert( + vectorSizeHints.empty() || + vectorSizeHints.size() == iterationDomain.size() && + "vector size hints should be empty or equal to the number of loops"); auto getStaticValue = [](Value v) -> int64_t { IntegerAttr attr; if (!matchPattern(v, m_Constant(&attr))) return ShapedType::kDynamicSize; @@ -281,7 +316,8 @@ distributedLoopUbs(numPartitionedLoops, ShapedType::kDynamicSize), minDistributedLoopTileSizes(numPartitionedLoops, 1), maxDistributedLoopTileSizes(numPartitionedLoops, - defaultWorkgroupTileSize); + defaultWorkgroupTileSize), + distributedVectorSizeHints(numPartitionedLoops, 1); // Find the bounds of the partitionable loops unsigned index = 0; for (auto range : llvm::enumerate(iterationDomain)) { @@ -291,13 +327,15 @@ maxDistributedLoopTileSizes[index] = maxTileSizes[range.index()]; distributedLoopLbs[index] = lbs[range.index()]; distributedLoopUbs[index] = ubs[range.index()]; + distributedVectorSizeHints[index] = + vectorSizeHints.empty() ? 1 : vectorSizeHints[range.index()]; index++; } SmallVector<int64_t> distributedTileSizes = - getDefaultDistributedLoopTileSizes(distributedLoopLbs, distributedLoopUbs, - minDistributedLoopTileSizes, - maxDistributedLoopTileSizes); + getDefaultDistributedLoopTileSizes( + distributedLoopLbs, distributedLoopUbs, minDistributedLoopTileSizes, + maxDistributedLoopTileSizes, distributedVectorSizeHints); SmallVector<int64_t> distributedLevelTileSizes(iterationDomain.size(), 0); for (auto loopID : llvm::enumerate(partitionableLoops)) { distributedLevelTileSizes[loopID.value()] = @@ -827,7 +865,10 @@ unsigned numLoops = convOp.getNumLoops(); SmallVector<int64_t> minTileSizes(numLoops, 1); SmallVector<int64_t> maxTileSizes(numLoops, defaultWorkgroupTileSize); + SmallVector<int64_t> vectorSizeHints(numLoops, 1); + // Give the vector size hint on OC. + vectorSizeHints[3] = vectorSize; // Set the flow level tiling to the default. OpBuilder builder(convOp.getContext()); builder.setInsertionPoint(convOp); @@ -837,7 +878,7 @@ cast<IREE::Flow::PartitionableLoopsInterface>(convOp.getOperation()); SmallVector<int64_t> flowTileSizes = getDefaultDistributedLevelTileSizes( iterationDomain, partitionableLoopsInterfaceOp, minTileSizes, - maxTileSizes); + maxTileSizes, vectorSizeHints); // Shapes of N, OH, OW, OC, KH, KW, (IC) Optional<SmallVector<int64_t, 4>> shapes = convOp.getStaticLoopRanges();
diff --git a/iree/compiler/Codegen/LLVMCPU/test/materialize_x86_64_launch_configuration.mlir b/iree/compiler/Codegen/LLVMCPU/test/materialize_x86_64_launch_configuration.mlir index b8fe472..ec807f3 100644 --- a/iree/compiler/Codegen/LLVMCPU/test/materialize_x86_64_launch_configuration.mlir +++ b/iree/compiler/Codegen/LLVMCPU/test/materialize_x86_64_launch_configuration.mlir
@@ -665,7 +665,7 @@ hal.executable private @depthwise_conv_static { hal.executable.variant public @system_elf_x86_64, target = <"llvm", "system-elf-x86_64", { data_layout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128", - native_vector_size = 16 : index, + native_vector_size = 64 : index, target_triple = "x86_64-unknown-linux-gnu" }> { hal.executable.entry_point public @depthwise_conv_static layout(#executable_layout) @@ -673,27 +673,27 @@ func.func @depthwise_conv_static() { %cst = arith.constant 0.0 : f32 %input_binding = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) - : !flow.dispatch.tensor<readonly:1x161x161x96xf32> + : !flow.dispatch.tensor<readonly:1x161x161x240xf32> %filter_binding = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) - : !flow.dispatch.tensor<readonly:3x3x96xf32> + : !flow.dispatch.tensor<readonly:3x3x240xf32> %result_binding = hal.interface.binding.subspan set(0) binding(2) type(storage_buffer) - : !flow.dispatch.tensor<writeonly:1x80x80x96xf32> - %input = flow.dispatch.tensor.load %input_binding, offsets = [0, 0, 0, 0], sizes = [1, 161, 161, 96], strides = [1, 1, 1, 1] - : !flow.dispatch.tensor<readonly:1x161x161x96xf32> -> tensor<1x161x161x96xf32> - %filter = flow.dispatch.tensor.load %filter_binding, offsets = [0, 0, 0], sizes = [3, 3, 96], strides = [1, 1, 1] - : !flow.dispatch.tensor<readonly:3x3x96xf32> -> tensor<3x3x96xf32> - %init = linalg.init_tensor [1, 80, 80, 96] : tensor<1x80x80x96xf32> - %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x80x80x96xf32>) -> tensor<1x80x80x96xf32> + : !flow.dispatch.tensor<writeonly:1x80x80x240xf32> + %input = flow.dispatch.tensor.load %input_binding, offsets = [0, 0, 0, 0], sizes = [1, 161, 161, 240], strides = [1, 1, 1, 1] + : !flow.dispatch.tensor<readonly:1x161x161x240xf32> -> tensor<1x161x161x240xf32> + %filter = flow.dispatch.tensor.load %filter_binding, offsets = [0, 0, 0], sizes = [3, 3, 240], strides = [1, 1, 1] + : !flow.dispatch.tensor<readonly:3x3x240xf32> -> tensor<3x3x240xf32> + %init = linalg.init_tensor [1, 80, 80, 240] : tensor<1x80x80x240xf32> + %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1x80x80x240xf32>) -> tensor<1x80x80x240xf32> %conv = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} - ins(%input, %filter : tensor<1x161x161x96xf32>, tensor<3x3x96xf32>) outs(%fill : tensor<1x80x80x96xf32>) -> tensor<1x80x80x96xf32> - flow.dispatch.tensor.store %conv, %result_binding, offsets = [0, 0, 0, 0], sizes = [1, 80, 80, 96], strides = [1, 1, 1, 1] - : tensor<1x80x80x96xf32> -> !flow.dispatch.tensor<writeonly:1x80x80x96xf32> + ins(%input, %filter : tensor<1x161x161x240xf32>, tensor<3x3x240xf32>) outs(%fill : tensor<1x80x80x240xf32>) -> tensor<1x80x80x240xf32> + flow.dispatch.tensor.store %conv, %result_binding, offsets = [0, 0, 0, 0], sizes = [1, 80, 80, 240], strides = [1, 1, 1, 1] + : tensor<1x80x80x240xf32> -> !flow.dispatch.tensor<writeonly:1x80x80x240xf32> return } } } } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[0, 20, 40, 48, 0, 0], [1, 1, 8, 8, 0, 0], [0, 0, 0, 0, 1, 3]]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_codegen.lowering_config<tile_sizes = {{\[}}[0, 40, 40, 48, 0, 0], [1, 1, 8, 16, 0, 0], [0, 0, 0, 0, 1, 3]]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<CPUConvTileAndDecomposeExpert> // CHECK: hal.executable.entry_point public @depthwise_conv_static // CHECK-SAME: translation_info = #[[TRANSLATION]]