[LLVMCPU] Align scalable distribution tiles to base * vscaleMax (#24849) Distribution tiles for scalable pack/unpack inner dims are aligned to the largest possible inner tile (`base * vscaleMax`) to ensure correctness and enable efficient tiling. Assisted-by: Claude Code --- <sub>Stack created with <a href="https://github.com/github/gh-stack">GitHub Stacks CLI</a> • <a href="https://gh.io/stacks-feedback">Give Feedback 💬</a></sub> --------- Signed-off-by: Ege Beysel <beyselege@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp index a1896ec..2248a21 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/KernelDispatch.cpp
@@ -617,6 +617,15 @@ return llvm::PowerOf2Ceil(size); } +/// Returns the maximum vscale recorded on the target, or 0 if no `vscale_range` +/// is set. +static int64_t getVscaleMax(IREE::HAL::ExecutableTargetAttr targetAttr) { + if (auto range = getVscaleRange(targetAttr)) { + return static_cast<int64_t>(range->vscaleMax); + } + return 0; +} + /// Computes the maximum tile size that can be used to distribute a dimension /// based on its number of iterations and the native vector size used of the /// target. The resulting tile size will be a multiple of the provided vector @@ -2039,12 +2048,14 @@ int srcRank = op.getSourceRank(); SmallVector<int64_t> innerTiles = op.getStaticTiles(); + IREE::Codegen::ScalableTileFlags scalableFlags(innerTiles.size(), false); // Try to infer scalable tile sizes. This is a no-op in case of static inner // tiles or if dynamic tile sizes are found, but scalable tile sizes cannot be // inferred. if (auto sizesAndScalableFlags = getScalableTileSizesAndFlags(op.getMixedTiles())) { innerTiles = sizesAndScalableFlags->first; + scalableFlags = sizesAndScalableFlags->second; } ArrayRef<int64_t> dimPos = op.getInnerDimsPos(); int64_t vectorSize = getVectorSize(entryPointFn, op.getSourceType()); @@ -2077,6 +2088,17 @@ pipelineConfig = getPipelineConfWithDecompositionAttr(op.getContext()); } + // For scalable inner tiles, the distribution tile on the packed dim must be a + // multiple of the largest possible inner tile (base * vscaleMax). + int64_t vscaleMax = getVscaleMax(target); + for (auto [pos, size, scalable] : + llvm::zip_equal(dimPos, innerTiles, scalableFlags)) { + if (!scalable || distTileSizes[pos] == 0 || ShapedType::isDynamic(size)) { + continue; + } + distTileSizes[pos] = llvm::alignTo(distTileSizes[pos], size * vscaleMax); + } + SmallVector<int64_t> vecTileSizes = getPackVectorTileSizes(entryPointFn, op); LoweringConfigGenerator generator(op); generator.setDistributionTileSizes(distTileSizes); @@ -2107,20 +2129,20 @@ innerTiles = sizesAndScalableFlags->first; scalableFlags = sizesAndScalableFlags->second; } - // Fixup for making distTileSizes be multiple of inner_tile_sizes. - // In case of scalable tile sizes, we align the distribution tile size with - // the static constant of the scalable tile size and round up to the next - // power of 2. Since vscale is a power of 2, this makes sure - // that the selected distribution size is divisible by or less than the - // effective scalable inner tile size. + // Align distribution tile sizes to inner tile sizes. For + // scalable inner tiles, align to the largest possible inner tile + // (base * vscaleMax) so the distribution tile stays a whole multiple of the + // runtime inner tile for any vscale. ArrayRef<int64_t> dimPos = op.getInnerDimsPos(); + int64_t vscaleMax = + getVscaleMax(IREE::HAL::ExecutableTargetAttr::lookup(entryPointFn)); for (auto [pos, size, scalable] : llvm::zip_equal(dimPos, innerTiles, scalableFlags)) { if (distTileSizes[pos] == 0 || ShapedType::isDynamic(size)) { continue; } - int64_t alignedTileSize = llvm::alignTo(distTileSizes[pos], size); - distTileSizes[pos] = roundUpToPow2(alignedTileSize, scalable); + int64_t multiplier = scalable ? vscaleMax : 1; + distTileSizes[pos] = llvm::alignTo(distTileSizes[pos], size * multiplier); } SmallVector<int64_t> vecTileSizes(op.getDestRank(), 1); @@ -3605,6 +3627,8 @@ void MultiLoweringConfigGenerator::adjustTileSizesForRootOp() { ArrayRef<int64_t> rootOpGlobalDims = dimTracker.getAllGlobalDimIdx(rootOperation); + int64_t vscaleMax = + getVscaleMax(IREE::HAL::ExecutableTargetAttr::lookup(rootOperation)); auto adjust = [&](Operation *op, ArrayRef<int64_t> vecTileSize, IREE::CPU::TilingLevel level, llvm::function_ref<int64_t(int64_t, int64_t)> updater) { @@ -3643,11 +3667,22 @@ continue; } // For pack op, align the distribution tile size and overwrite the - // vector parallel tile size and scalable flag. - adjust(op, vecTileSize, IREE::CPU::TilingLevel::DistributionTiles, align); + // vector parallel tile size and scalable flag. Distribution tiles for + // scalable inner tiles are aligned to their largest runtime size (base * + // vscaleMax). + SmallVector<bool> packScalableFlags = nonRootOpScalableFlags.lookup(op); + SmallVector<int64_t> distAlignTiles(vecTileSize.begin(), + vecTileSize.end()); + for (auto [i, scalable] : llvm::enumerate(packScalableFlags)) { + if (scalable && i < distAlignTiles.size() && distAlignTiles[i] > 0) { + distAlignTiles[i] *= vscaleMax; + } + } + adjust(op, distAlignTiles, IREE::CPU::TilingLevel::DistributionTiles, + align); adjust(op, vecTileSize, IREE::CPU::TilingLevel::VectorCommonParallelTiles, overwrite); - adjustScalableFlags(op, nonRootOpScalableFlags.lookup(op), + adjustScalableFlags(op, packScalableFlags, IREE::CPU::TilingLevel::VectorCommonParallelTiles, overwrite); } else if (auto unpackOp = dyn_cast<linalg::UnPackOp>(op)) { @@ -4066,10 +4101,11 @@ LDBG() << "The tile sizes for each dimension should be aligned to " << alignedSizes; - // Fixup for making tileSizes be multiple of inner_tile_sizes. In case of - // scalable inner tiles, we align the distribution tile sizes with the static - // constant of the scalable inner tile size and round up to the next power of - // 2 to ensure alignment. + // Fixup for making tileSizes be multiple of inner_tile_sizes. Scalable + // distribution tiles are aligned to the largest possible inner tile + // (base * vscaleMax). + int64_t vscaleMax = + getVscaleMax(IREE::HAL::ExecutableTargetAttr::lookup(entryPointFn)); SmallVector<IREE::CPU::LoweringConfigLevelInfo> tilingInfo = loweringConfig.getAvailableTilingInfo(); for (IREE::CPU::LoweringConfigLevelInfo &info : tilingInfo) { @@ -4078,12 +4114,12 @@ if (tileSizes[idx] == 0) { continue; } - int64_t alignedTileSize = - llvm::alignTo(tileSizes[idx], alignedSizes[idx]); - tileSizes[idx] = roundUpToPow2( - alignedTileSize, - vecParallelScalableTileFlags[idx] && - info.level == IREE::CPU::TilingLevel::DistributionTiles); + int multiplier = (vecParallelScalableTileFlags[idx] && + info.level == IREE::CPU::TilingLevel::DistributionTiles) + ? vscaleMax + : 1; + tileSizes[idx] = + llvm::alignTo(tileSizes[idx], alignedSizes[idx] * multiplier); } // Fixup for the scalable tile flags. if (info.level == IREE::CPU::TilingLevel::VectorCommonParallelTiles) {
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir index 27fd101..c22c451 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_aarch64_sve_lowering_strategy.mlir
@@ -181,7 +181,7 @@ %unpack = linalg.unpack %arg0 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, %c8_vscale] into %init : tensor<128x10x?x8x?xf32> -> tensor<128x80x320xf32> return %unpack : tensor<128x80x320xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 40, 64], vector_common_parallel = [1, 8, [8]]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 40, 128], vector_common_parallel = [1, 8, [8]]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling> // // For SVE, we do not decompose unpacks. @@ -203,7 +203,7 @@ %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [32, %c8_vscale] into %init : tensor<?x?x32x?xi32> -> tensor<?x?xi32> return %unpack : tensor<?x?xi32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [32, [8]]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 128], vector_common_parallel = [32, [8]]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling> // // For SVE, we do not decompose unpacks. @@ -227,7 +227,7 @@ %unpack = linalg.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [32, %c8_vscale] into %init : tensor<?x?x32x?xf32> -> tensor<?x?xf32> return %unpack : tensor<?x?xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = {{\[\[}}8], 32]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [128, 64], vector_common_parallel = {{\[\[}}8], 32]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling> // CHECK: func.func @unpack_transposed // CHECK-SAME: translation_info = #[[TRANSLATION]] @@ -274,7 +274,7 @@ return %generic : tensor<128x320x80xf32> } // CHECK-DAG: #[[CONFIG_UNPACK:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 8, [8]]> -// CHECK-DAG: #[[CONFIG_GENERIC:.+]] = #iree_cpu.lowering_config<distribution = [64, 64, 40], vector_common_parallel = [1, [8], 8]> +// CHECK-DAG: #[[CONFIG_GENERIC:.+]] = #iree_cpu.lowering_config<distribution = [64, 128, 40], vector_common_parallel = [1, [8], 8]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert> // CHECK: func.func @unpack_with_generic // CHECK-SAME: translation_info = #[[TRANSLATION]] @@ -298,7 +298,7 @@ %pack = linalg.pack %arg0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [16, %c16_vscale] into %empty : tensor<20x48xf32> -> tensor<2x?x16x?xf32> return %pack : tensor<2x?x16x?xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [1, 64], vector_common_parallel = [1, 1]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [1, 256], vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling>> // CHECK: func.func @pack( // CHECK-SAME: translation_info = #[[TRANSLATION]] @@ -326,7 +326,7 @@ %pack = linalg.pack %filled inner_dims_pos = [0, 1] inner_tiles = [8, %c16_vscale] into %dest : tensor<128x384xf32> -> tensor<16x?x8x?xf32> return %pack : tensor<16x?x8x?xf32> } -// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [8, [16]]> +// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 256], vector_common_parallel = [8, [16]]> // CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>> // CHECK: func.func @elem_pack( @@ -405,7 +405,7 @@ %pack = linalg.pack %0 padding_value(%cst : f32) outer_dims_perm = [1, 0] inner_dims_pos = [1, 0] inner_tiles = [%c16_vscale, 8] into %empty1 : tensor<768x30522xf32> -> tensor<?x96x?x8xf32> return %pack : tensor<?x96x?x8xf32> } -// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [8, [16]]> +// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 256], vector_common_parallel = [8, [16]]> // CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>> // CHECK: func.func @transpose_pack( @@ -442,7 +442,7 @@ return %pack : tensor<?x24x16x?xf32> } // CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [16, [16]]> -// CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<distribution = [32, 32, 0], vector_common_parallel = [16, [16], 0], vector_reduction = [0, 0, 4]> +// CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<distribution = [32, 256, 0], vector_common_parallel = [16, [16], 0], vector_reduction = [0, 0, 4]> // CHECK-DAG: #[[CONFIG3:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>> // CHECK: func.func @reduction_pack(
diff --git a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_v_lowering_strategy.mlir b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_v_lowering_strategy.mlir index 1abfc93..c5ff593 100644 --- a/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_v_lowering_strategy.mlir +++ b/compiler/src/iree/compiler/Codegen/LLVMCPU/test/select_riscv_v_lowering_strategy.mlir
@@ -69,7 +69,7 @@ %unpack = linalg.unpack %arg0 outer_dims_perm = [0, 1, 2] inner_dims_pos = [1, 2] inner_tiles = [8, %c8_vscale] into %init : tensor<128x10x?x8x?xf32> -> tensor<128x80x320xf32> return %unpack : tensor<128x80x320xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 40, 64], vector_common_parallel = [1, 8, [8]]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 40, 128], vector_common_parallel = [1, 8, [8]]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling>> // // For RVV, we do not decompose unpacks. @@ -91,7 +91,7 @@ %unpack = linalg.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [32, %c8_vscale] into %init : tensor<?x?x32x?xf32> -> tensor<?x?xf32> return %unpack : tensor<?x?xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [32, [8]]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [64, 128], vector_common_parallel = [32, [8]]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling>> // // For RVV, we do not decompose unpacks. @@ -116,7 +116,7 @@ %pack = linalg.pack %arg0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [16, %c16_vscale] into %empty : tensor<20x48xf32> -> tensor<2x?x16x?xf32> return %pack : tensor<2x?x16x?xf32> } -// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [1, 64], vector_common_parallel = [1, 1]> +// CHECK-DAG: #[[CONFIG:.+]] = #iree_cpu.lowering_config<distribution = [1, 256], vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DataTiling>> // CHECK: func.func @pack( // CHECK-SAME: translation_info = #[[TRANSLATION]] @@ -144,7 +144,7 @@ %pack = linalg.pack %filled inner_dims_pos = [0, 1] inner_tiles = [8, %c16_vscale] into %dest : tensor<128x384xf32> -> tensor<16x?x8x?xf32> return %pack : tensor<16x?x8x?xf32> } -// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 64], vector_common_parallel = [8, [16]]> +// CHECK-DAG: #[[CONFIG1:.+]] = #iree_cpu.lowering_config<distribution = [64, 256], vector_common_parallel = [8, [16]]> // CHECK-DAG: #[[CONFIG2:.+]] = #iree_cpu.lowering_config<vector_common_parallel = [1, 1]> // CHECK-DAG: #[[TRANSLATION:.+]] = #iree_codegen.translation_info<pipeline = #iree_cpu.pipeline<DoubleTilingExpert>> // CHECK: func.func @elem_pack(