[spirv] Add conv CodeGen configuration for Adreno GPUs (#7001) Similar to matmul, this enables tiling and vectorizing conv ops for Adreno GPUs. The goal is to eventually promote what's here as the default for all targets, including Mali GPUs, so that we can remove the hardcoded known good configurations. But for now starting with Adreno to make sure they work fine. Along the way, moved the "fourth-level" tiling into the pattern itself. It's really fundamentally different than the others. We don't have four levels of compute hierarchies on GPU to map to. We just used it as a way to materialize the loops around filter height and width. It's required for applicability and correctness for the pattern itself and is always the same, not tunable knobs that change according to the input.
diff --git a/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp index 39a3bbc..3da6a17 100644 --- a/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp +++ b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp
@@ -59,7 +59,8 @@ // Deduce the configuration for the N dimension. Start with the best workgroup // X size, and reduce by a factor of two each time. for (int64_t x = bestX; x >= 2; x >>= 1) { - // Handle 4 elements per thread. We need this for vectorized load. + // Handle 4 elements per thread for the innermost dimension. We need this + // for vectorized load. int64_t chosenTileSize = 4; if (dimN % (x * chosenTileSize) == 0) { workgroupSize[0] = x; @@ -116,6 +117,112 @@ } //===----------------------------------------------------------------------===// +// Convolution +//===----------------------------------------------------------------------===// + +/// Sets CodeGen configurations via attributes to the given convolution +/// `linalgOp` by trying to achieve the given `bestTilingFactor`, which is how +/// many scalar elements each thread should handle. +static LogicalResult setOpConfig(linalg::LinalgOp linalgOp, + int64_t bestTilingFactor) { + ArrayRef<int64_t> inputShape = getUntiledShape(linalgOp.inputs()[0]); + ArrayRef<int64_t> outputShape = getUntiledResultShape(linalgOp, 0); + if (llvm::any_of(inputShape, ShapedType::isDynamic)) return success(); + if (llvm::any_of(outputShape, ShapedType::isDynamic)) return success(); + + int64_t ic = inputShape[3]; + int64_t oh = outputShape[1], ow = outputShape[2], oc = outputShape[3]; + + // The conversion pipeline requires the input channel dimension to be some + // multipler of four, or less than four. + if (!(ic % 4 == 0 || ic < 4)) return success(); + + // The core idea is to distribute the convolution OH/OW/OC dimension to the + // workgroup Z/Y/X dimension, with each thread in a workgroup handling + // multiple vector elements. We try to 1) utilize all threads in a subgroup, + // and 2) handle an optimal tile size along each dimension. + + int64_t residualThreads = 64; + int64_t residualTilingFactor = bestTilingFactor; + + SmallVector<int64_t, 3> workgroupSize(3, 1); // (X, Y, Z) + SmallVector<int64_t, 4> workgroupTileSizes(4, 0); // (N, OH, OW, OC) + SmallVector<int64_t, 4> invocationTileSizes(4, 0); // (N, OH, OW, OC) + + // Deduce the configuration for the OC dimension. + for (int64_t x = residualThreads; x >= 2; x >>= 1) { + // Handle 4 elements per thread for the innermost dimension. We need this + // for vectorized load. + int64_t chosenTileSize = 4; + if (oc % (x * chosenTileSize) == 0) { + workgroupSize[0] = x; + workgroupTileSizes[3] = x * chosenTileSize; + invocationTileSizes[3] = chosenTileSize; + residualThreads /= x; + residualTilingFactor /= chosenTileSize; + break; + } + } + if (workgroupTileSizes[3] == 0) return success(); + + // Deduce the configruation for the OW and OH dimension. Try to make them even + // if possible given we typically have images with the same height and width. + if (ow == oh && residualThreads != 1 && + llvm::Log2_64(residualThreads) % 2 == 0) { + int64_t yz = 1 << (llvm::Log2_64(residualThreads) / 2); + int64_t chosenTileSize = 1 << (llvm::Log2_64(residualTilingFactor) / 2); + while (ow % chosenTileSize != 0) chosenTileSize >>= 1; + workgroupSize[1] = workgroupSize[2] = yz; + workgroupTileSizes[2] = workgroupTileSizes[1] = yz * chosenTileSize; + invocationTileSizes[2] = invocationTileSizes[1] = chosenTileSize; + } else { + auto decideOneDim = [&](int64_t inputDim, int64_t &wgDimSize, + int64_t &wgTileSize, int64_t &invoTileSize) { + for (int64_t dim = residualThreads; dim >= 1; dim >>= 1) { + int64_t chosenTileSize = 0; + for (int64_t t = residualTilingFactor; t >= 1; t >>= 1) { + if (inputDim % (dim * t) == 0) { + chosenTileSize = t; + break; + } + } + if (chosenTileSize) { + wgDimSize = dim; + wgTileSize = dim * chosenTileSize; + invoTileSize = chosenTileSize; + residualThreads /= dim; + residualTilingFactor /= chosenTileSize; + return true; + } + } + return false; + }; + + if (!decideOneDim(ow, workgroupSize[1], workgroupTileSizes[2], + invocationTileSizes[2]) || + !decideOneDim(oh, workgroupSize[2], workgroupTileSizes[1], + invocationTileSizes[1])) { + return success(); + } + } + + auto pipeline = IREE::HAL::DispatchLoweringPassPipeline::SPIRVVectorize; + TileSizesListType tileSizes; + tileSizes.push_back(workgroupTileSizes); + tileSizes.emplace_back(); // Subgroup level + tileSizes.push_back(invocationTileSizes); + + auto funcOp = linalgOp->getParentOfType<FuncOp>(); + if (failed(setOpConfigAndEntryPointFnTranslation( + funcOp, linalgOp, tileSizes, {}, pipeline, workgroupSize))) { + return failure(); + } + return defineConvWorkgroupCountRegion( + linalgOp, llvm::makeArrayRef(outputShape).drop_front(), + llvm::makeArrayRef(workgroupTileSizes).drop_front()); +} + +//===----------------------------------------------------------------------===// // Entry Point //===----------------------------------------------------------------------===// @@ -124,6 +231,10 @@ return TypeSwitch<Operation *, LogicalResult>(rootOp) .Case<linalg::BatchMatmulOp, linalg::MatmulOp>( [](auto op) { return setOpConfig(op); }) + .Case<linalg::Conv2DNhwcHwcfOp>( + [](auto op) { return setOpConfig(op, 32); }) + .Case<linalg::DepthwiseConv2DNhwOp>( + [](auto op) { return setOpConfig(op, 16); }) .Default([](Operation *) { return success(); }); }
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp index d7e72f3..7ad7085 100644 --- a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp +++ b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -67,6 +67,24 @@ return defineWorkgroupCountRegion(builder, funcOp, numWorkgroupsFn); } +namespace detail { +LogicalResult defineConvWorkgroupCountRegion( + Operation *op, ArrayRef<int64_t> outputShape, + ArrayRef<int64_t> workgroupTileSizes) { + auto numWorkgroupsFn = [&](OpBuilder &b, Location loc, std::array<Value, 3>) { + std::array<Value, 3> xyz; + for (unsigned i = 0; i < 3; ++i) { + int64_t count = outputShape[i] / workgroupTileSizes[i]; + xyz[2 - i] = b.create<ConstantIndexOp>(loc, count); + } + return xyz; + }; + OpBuilder builder(op->getContext()); + return defineWorkgroupCountRegion(builder, op->getParentOfType<FuncOp>(), + numWorkgroupsFn); +} +} // namespace detail + //===----------------------------------------------------------------------===// // Matmul Default Configuration //===----------------------------------------------------------------------===//
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.h b/iree/compiler/Codegen/SPIRV/KernelConfig.h index 61f2a13..c8641ce 100644 --- a/iree/compiler/Codegen/SPIRV/KernelConfig.h +++ b/iree/compiler/Codegen/SPIRV/KernelConfig.h
@@ -26,6 +26,14 @@ namespace detail { +/// Lets the entry point region to return fully static number of workgroups. +// This is needed for folding `affine.min` ops to expose static-shaped tiled +// convolution for vectorization. +// TODO(#5034): Use a proper way to prove tilability and fold `affine.min`s. +LogicalResult defineConvWorkgroupCountRegion( + Operation *op, ArrayRef<int64_t> outputShape, + ArrayRef<int64_t> workgroupTileSizes); + /// Sets CodeGen configuration for GPUs from a specific vendor. /// /// If the given `rootOp` has known good CodeGen configuration, attaches a
diff --git a/iree/compiler/Codegen/SPIRV/MaliConfig.cpp b/iree/compiler/Codegen/SPIRV/MaliConfig.cpp index ff404f0..b933aec 100644 --- a/iree/compiler/Codegen/SPIRV/MaliConfig.cpp +++ b/iree/compiler/Codegen/SPIRV/MaliConfig.cpp
@@ -226,10 +226,6 @@ /*output_height=*/tileSize[0] / workgroupSize[2], /*output_width=*/tileSize[1] / workgroupSize[1], /*output_channel=*/tileSize[2] / workgroupSize[0]}); - // Finally, for each invocation, we use tiling to generate loops to loop - // over the filter's height (step 1), width (step 1), and input channel - // (step 4) dimensions. - tileSizes.push_back({0, 0, 0, 0, 1, 1, 4}); auto funcOp = op->getParentOfType<FuncOp>(); if (failed(setOpConfigAndEntryPointFnTranslation( @@ -237,21 +233,8 @@ return failure(); } - // Let the entry point region to return fully static number of workgroups. - // This is needed for folding `affine.min` ops to expose static-shaped tiled - // convolution for vectorization. - // TODO(#5034): Use a proper way to prove tilability and fold `affine.min`s. - OpBuilder builder(op.getContext()); - auto numWorkgroupsFn = [&](OpBuilder &b, Location loc, - std::array<Value, 3>) { - std::array<Value, 3> xyz; - for (unsigned i = 0; i < 3; ++i) { - int64_t count = outputShape[i + 1] / tileSize[i]; - xyz[2 - i] = b.create<ConstantIndexOp>(loc, count); - } - return xyz; - }; - return defineWorkgroupCountRegion(builder, funcOp, numWorkgroupsFn); + return defineConvWorkgroupCountRegion( + op, llvm::makeArrayRef(outputShape).drop_front(), tileSize); } return success(); } @@ -302,9 +285,6 @@ /*output_height=*/tileSize[0] / workgroupSize[2], /*output_width=*/tileSize[1] / workgroupSize[1], /*output_channel=*/tileSize[2] / workgroupSize[0]}); - // Finally, for each invocation, we use tiling to generate loops to loop - // over the filter's height (step 1) and width (step 1) dimensions. - tileSizes.push_back({0, 0, 0, 0, 1, 1}); auto funcOp = op->getParentOfType<FuncOp>(); if (failed(setOpConfigAndEntryPointFnTranslation( @@ -312,21 +292,8 @@ return failure(); } - // Let the entry point region to return fully static number of workgroups. - // This is needed for folding `affine.min` ops to expose static-shaped tiled - // convolution for vectorization. - // TODO(#5034): Use a proper way to prove tilability and fold `affine.min`s. - OpBuilder builder(op.getContext()); - auto numWorkgroupsFn = [&](OpBuilder &b, Location loc, - std::array<Value, 3>) { - std::array<Value, 3> xyz; - for (unsigned i = 0; i < 3; ++i) { - int64_t count = outputShape[i + 1] / tileSize[i]; - xyz[2 - i] = b.create<ConstantIndexOp>(loc, count); - } - return xyz; - }; - return defineWorkgroupCountRegion(builder, funcOp, numWorkgroupsFn); + return defineConvWorkgroupCountRegion( + op, llvm::makeArrayRef(outputShape).drop_front(), tileSize); } return success(); }
diff --git a/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp b/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp index 8736927..8e76ee8 100644 --- a/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp +++ b/iree/compiler/Codegen/SPIRV/SPIRVTileAndDistribute.cpp
@@ -238,12 +238,16 @@ MLIRContext *context, RewritePatternSet &patterns, linalg::LinalgTransformationFilter marker) { auto getTileSizeFn = [&](OpBuilder &builder, Operation *op) { - SmallVector<Value, 4> tileSizes; - SmallVector<int64_t, 4> fourthLevel = getTileSizes(op, 3); - tileSizes.reserve(fourthLevel.size()); + SmallVector<int64_t, 4> sizes; + if (isa<linalg::Conv2DNhwcHwcfOp>(op)) { + sizes.assign({0, 0, 0, 0, 1, 1, 4}); + } else if (isa<linalg::DepthwiseConv2DNhwOp>(op)) { + sizes.assign({0, 0, 0, 0, 1, 1}); + } Location loc = op->getLoc(); - for (int64_t size : fourthLevel) { + SmallVector<Value, 4> tileSizes; + for (int64_t size : sizes) { tileSizes.push_back(builder.create<ConstantIndexOp>(loc, size)); } return tileSizes; @@ -254,8 +258,7 @@ .setTileSizeComputationFunction(getTileSizeFn); patterns.insert<linalg::LinalgTilingPattern<linalg::Conv2DNhwcHwcfOp>, - linalg::LinalgTilingPattern<linalg::DepthwiseConv2DNhwOp>, - linalg::LinalgTilingPattern<linalg::DepthwiseConv2DNhwcOp>>( + linalg::LinalgTilingPattern<linalg::DepthwiseConv2DNhwOp>>( context, tilingOptions, marker); }
diff --git a/iree/compiler/Codegen/SPIRV/test/BUILD b/iree/compiler/Codegen/SPIRV/test/BUILD index 330e6a3..f4dae26 100644 --- a/iree/compiler/Codegen/SPIRV/test/BUILD +++ b/iree/compiler/Codegen/SPIRV/test/BUILD
@@ -19,6 +19,7 @@ name = "lit", srcs = enforce_glob( [ + "config_adreno_conv.mlir", "config_adreno_matmul.mlir", "config_linalg_ext_ops.mlir", "convert_to_spirv.mlir",
diff --git a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt index ef888fb..22f45ca 100644 --- a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt +++ b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
@@ -14,6 +14,7 @@ NAME lit SRCS + "config_adreno_conv.mlir" "config_adreno_matmul.mlir" "config_linalg_ext_ops.mlir" "convert_to_spirv.mlir"
diff --git a/iree/compiler/Codegen/SPIRV/test/config_adreno_conv.mlir b/iree/compiler/Codegen/SPIRV/test/config_adreno_conv.mlir new file mode 100644 index 0000000..31b68c8 --- /dev/null +++ b/iree/compiler/Codegen/SPIRV/test/config_adreno_conv.mlir
@@ -0,0 +1,441 @@ +// RUN: iree-opt -split-input-file -mlir-print-local-scope -pass-pipeline='hal.executable(hal.executable.variant(iree-spirv-lower-executable-target-pass{test-lowering-configuration=true}))' %s | IreeFileCheck %s + +// Conv - large OC - distribute to only one workgroup dimension. + +hal.executable @conv_112x112x512 { + hal.interface public @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant public @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", { + spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, { + max_compute_shared_memory_size = 32768 : i32, + max_compute_workgroup_invocations = 1024 : i32, + max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>, + subgroup_size = 64 : i32}> + }> { + hal.executable.entry_point public @conv_112x112x512 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @conv_112x112x512() { + %c0 = constant 0 : index + %c512 = constant 512 : index + %c112 = constant 112 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1x225x225x3xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:3x3x3x512xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1x112x112x512xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + %workgroup_size_z = hal.interface.workgroup.size[2] : index + %workgroup_id_x = hal.interface.workgroup.id[0] : index + %workgroup_count_x = hal.interface.workgroup.count[0] : index + %workgroup_id_y = hal.interface.workgroup.id[1] : index + %workgroup_count_y = hal.interface.workgroup.count[1] : index + %workgroup_id_z = hal.interface.workgroup.id[2] : index + %workgroup_count_z = hal.interface.workgroup.count[2] : index + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_z, %workgroup_size_z] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_z, %workgroup_size_z] + scf.for %arg0 = %3 to %c112 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg1 = %5 to %c112 step %6 { + %7 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %8 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg2 = %7 to %c512 step %8 { + %9 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg0) + %10 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 225)>(%arg0)[%workgroup_size_z] + %11 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg1) + %12 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 225)>(%arg1)[%workgroup_size_y] + %13 = flow.dispatch.tensor.load %0, offsets = [0, %9, %11, 0], sizes = [1, %10, %12, 3], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:1x225x225x3xf32> -> tensor<1x?x?x3xf32> + %14 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 512)>(%arg2)[%workgroup_size_x] + %15 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, %arg2], sizes = [3, 3, 3, %14], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:3x3x3x512xf32> -> tensor<3x3x3x?xf32> + %16 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 112)>(%arg0)[%workgroup_size_z] + %17 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 112)>(%arg1)[%workgroup_size_y] + %18 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 512)>(%arg2)[%workgroup_size_x] + %19 = affine.min affine_map<(d0)[s0] -> (-d0 + 112, s0)>(%arg0)[%workgroup_size_z] + %20 = affine.min affine_map<(d0)[s0] -> (-d0 + 112, s0)>(%arg1)[%workgroup_size_y] + %21 = affine.min affine_map<(d0)[s0] -> (-d0 + 512, s0)>(%arg2)[%workgroup_size_x] + %22 = linalg.init_tensor [1, %19, %20, %21] : tensor<1x?x?x?xf32> + %23 = linalg.fill(%cst, %22) : f32, tensor<1x?x?x?xf32> -> tensor<1x?x?x?xf32> + %24 = linalg.conv_2d_nhwc_hwcf {__internal_linalg_transform__ = "workgroup", dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%13, %15 : tensor<1x?x?x3xf32>, tensor<3x3x3x?xf32>) outs(%23 : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32> + flow.dispatch.tensor.store %24, %2, offsets = [0, %arg0, %arg1, %arg2], sizes = [1, %16, %17, %18], strides = [1, 1, 1, 1] : tensor<1x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:1x112x112x512xf32> + } + } + } + return + } + hal.interface private @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @conv_112x112x512 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [256, 8, 1]} +// CHECK-SAME: workgroup_size = [64 : index, 1 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[C2:.+]] = constant 2 : index +// CHECK-NEXT: %[[C14:.+]] = constant 14 : index +// CHECK-NEXT: %[[C112:.+]] = constant 112 : index +// CHECK-NEXT: hal.return %[[C2]], %[[C14]], %[[C112]] + +// CHECK: func @conv_112x112x512() +// CHECK: linalg.conv_2d_nhwc_hwcf +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[0, 1, 8, 256], [], [0, 1, 8, 4]]} + +// ----- + +// Conv - medium OC/OW/OH - distribute to two workgroup dimensions. + +hal.executable @conv_112x112x32 { + hal.interface public @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant public @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", { + spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, { + max_compute_shared_memory_size = 32768 : i32, + max_compute_workgroup_invocations = 1024 : i32, + max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>, + subgroup_size = 64 : i32}> + }> { + hal.executable.entry_point public @conv_112x112x32 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @conv_112x112x32() { + %c0 = constant 0 : index + %c32 = constant 32 : index + %c112 = constant 112 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1x225x225x3xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:3x3x3x32xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1x112x112x32xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + %workgroup_size_z = hal.interface.workgroup.size[2] : index + %workgroup_id_x = hal.interface.workgroup.id[0] : index + %workgroup_count_x = hal.interface.workgroup.count[0] : index + %workgroup_id_y = hal.interface.workgroup.id[1] : index + %workgroup_count_y = hal.interface.workgroup.count[1] : index + %workgroup_id_z = hal.interface.workgroup.id[2] : index + %workgroup_count_z = hal.interface.workgroup.count[2] : index + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_z, %workgroup_size_z] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_z, %workgroup_size_z] + scf.for %arg0 = %3 to %c112 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg1 = %5 to %c112 step %6 { + %7 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %8 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg2 = %7 to %c32 step %8 { + %9 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg0) + %10 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 225)>(%arg0)[%workgroup_size_z] + %11 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg1) + %12 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 225)>(%arg1)[%workgroup_size_y] + %13 = flow.dispatch.tensor.load %0, offsets = [0, %9, %11, 0], sizes = [1, %10, %12, 3], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:1x225x225x3xf32> -> tensor<1x?x?x3xf32> + %14 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 32)>(%arg2)[%workgroup_size_x] + %15 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, %arg2], sizes = [3, 3, 3, %14], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:3x3x3x32xf32> -> tensor<3x3x3x?xf32> + %16 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 112)>(%arg0)[%workgroup_size_z] + %17 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 112)>(%arg1)[%workgroup_size_y] + %18 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 32)>(%arg2)[%workgroup_size_x] + %19 = affine.min affine_map<(d0)[s0] -> (-d0 + 112, s0)>(%arg0)[%workgroup_size_z] + %20 = affine.min affine_map<(d0)[s0] -> (-d0 + 112, s0)>(%arg1)[%workgroup_size_y] + %21 = affine.min affine_map<(d0)[s0] -> (-d0 + 32, s0)>(%arg2)[%workgroup_size_x] + %22 = linalg.init_tensor [1, %19, %20, %21] : tensor<1x?x?x?xf32> + %23 = linalg.fill(%cst, %22) : f32, tensor<1x?x?x?xf32> -> tensor<1x?x?x?xf32> + %24 = linalg.conv_2d_nhwc_hwcf {__internal_linalg_transform__ = "workgroup", dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%13, %15 : tensor<1x?x?x3xf32>, tensor<3x3x3x?xf32>) outs(%23 : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32> + flow.dispatch.tensor.store %24, %2, offsets = [0, %arg0, %arg1, %arg2], sizes = [1, %16, %17, %18], strides = [1, 1, 1, 1] : tensor<1x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:1x112x112x32xf32> + } + } + } + return + } + hal.interface private @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @conv_112x112x32 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [32, 16, 4]} +// CHECK-SAME: workgroup_size = [8 : index, 8 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[C1:.+]] = constant 1 : index +// CHECK-NEXT: %[[C7:.+]] = constant 7 : index +// CHECK-NEXT: %[[C28:.+]] = constant 28 : index +// CHECK-NEXT: hal.return %[[C1]], %[[C7]], %[[C28]] + +// CHECK: func @conv_112x112x32() +// CHECK: linalg.conv_2d_nhwc_hwcf +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[0, 4, 16, 32], [], [0, 4, 2, 4]]} + +// ----- + +// Conv - small OC/OW/OH - distribute to all three workgroup dimensions. + +hal.executable @conv_16x16x16 { + hal.interface public @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant public @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", { + spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, { + max_compute_shared_memory_size = 32768 : i32, + max_compute_workgroup_invocations = 1024 : i32, + max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>, + subgroup_size = 64 : i32}> + }> { + hal.executable.entry_point public @conv_16x16x16 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @conv_16x16x16() { + %c0 = constant 0 : index + %c16 = constant 16 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1x33x33x3xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:3x3x3x16xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1x16x16x16xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + %workgroup_size_z = hal.interface.workgroup.size[2] : index + %workgroup_id_x = hal.interface.workgroup.id[0] : index + %workgroup_count_x = hal.interface.workgroup.count[0] : index + %workgroup_id_y = hal.interface.workgroup.id[1] : index + %workgroup_count_y = hal.interface.workgroup.count[1] : index + %workgroup_id_z = hal.interface.workgroup.id[2] : index + %workgroup_count_z = hal.interface.workgroup.count[2] : index + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_z, %workgroup_size_z] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_z, %workgroup_size_z] + scf.for %arg0 = %3 to %c16 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg1 = %5 to %c16 step %6 { + %7 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %8 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg2 = %7 to %c16 step %8 { + %9 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg0) + %10 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 33)>(%arg0)[%workgroup_size_z] + %11 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg1) + %12 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 33)>(%arg1)[%workgroup_size_y] + %13 = flow.dispatch.tensor.load %0, offsets = [0, %9, %11, 0], sizes = [1, %10, %12, 3], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:1x33x33x3xf32> -> tensor<1x?x?x3xf32> + %14 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 16)>(%arg2)[%workgroup_size_x] + %15 = flow.dispatch.tensor.load %1, offsets = [0, 0, 0, %arg2], sizes = [3, 3, 3, %14], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:3x3x3x16xf32> -> tensor<3x3x3x?xf32> + %16 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 16)>(%arg0)[%workgroup_size_z] + %17 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 16)>(%arg1)[%workgroup_size_y] + %18 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 16)>(%arg2)[%workgroup_size_x] + %19 = affine.min affine_map<(d0)[s0] -> (-d0 + 16, s0)>(%arg0)[%workgroup_size_z] + %20 = affine.min affine_map<(d0)[s0] -> (-d0 + 16, s0)>(%arg1)[%workgroup_size_y] + %21 = affine.min affine_map<(d0)[s0] -> (-d0 + 16, s0)>(%arg2)[%workgroup_size_x] + %22 = linalg.init_tensor [1, %19, %20, %21] : tensor<1x?x?x?xf32> + %23 = linalg.fill(%cst, %22) : f32, tensor<1x?x?x?xf32> -> tensor<1x?x?x?xf32> + %24 = linalg.conv_2d_nhwc_hwcf {__internal_linalg_transform__ = "workgroup", dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%13, %15 : tensor<1x?x?x3xf32>, tensor<3x3x3x?xf32>) outs(%23 : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32> + flow.dispatch.tensor.store %24, %2, offsets = [0, %arg0, %arg1, %arg2], sizes = [1, %16, %17, %18], strides = [1, 1, 1, 1] : tensor<1x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:1x16x16x16xf32> + } + } + } + return + } + hal.interface private @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @conv_16x16x16 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [16, 8, 8]} +// CHECK-SAME: workgroup_size = [4 : index, 4 : index, 4 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[C1:.+]] = constant 1 : index +// CHECK-NEXT: %[[C2:.+]] = constant 2 : index +// CHECK-NEXT: hal.return %[[C1]], %[[C2]], %[[C2]] + +// CHECK: func @conv_16x16x16() +// CHECK: linalg.conv_2d_nhwc_hwcf +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[0, 8, 8, 16], [], [0, 2, 2, 4]]} + +// ----- + +// Depthwise conv - small OC/OW/OH - distribute to all three workgroup dimensions. + +hal.executable @dwconv_28x28x144 { + hal.interface public @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant public @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", { + spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, { + max_compute_shared_memory_size = 32768 : i32, + max_compute_workgroup_invocations = 1024 : i32, + max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>, + subgroup_size = 64 : i32}> + }> { + hal.executable.entry_point public @dwconv_28x28x144 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @dwconv_28x28x144() { + %c0 = constant 0 : index + %c144 = constant 144 : index + %c28 = constant 28 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1x57x57x144xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:3x3x144xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1x28x28x144xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + %workgroup_size_z = hal.interface.workgroup.size[2] : index + %workgroup_id_x = hal.interface.workgroup.id[0] : index + %workgroup_count_x = hal.interface.workgroup.count[0] : index + %workgroup_id_y = hal.interface.workgroup.id[1] : index + %workgroup_count_y = hal.interface.workgroup.count[1] : index + %workgroup_id_z = hal.interface.workgroup.id[2] : index + %workgroup_count_z = hal.interface.workgroup.count[2] : index + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_z, %workgroup_size_z] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_z, %workgroup_size_z] + scf.for %arg0 = %3 to %c28 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg1 = %5 to %c28 step %6 { + %7 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %8 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg2 = %7 to %c144 step %8 { + %9 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg0) + %10 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 57)>(%arg0)[%workgroup_size_z] + %11 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg1) + %12 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 57)>(%arg1)[%workgroup_size_y] + %13 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 144)>(%arg2)[%workgroup_size_x] + %14 = flow.dispatch.tensor.load %0, offsets = [0, %9, %11, %arg2], sizes = [1, %10, %12, %13], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:1x57x57x144xf32> -> tensor<1x?x?x?xf32> + %15 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 144)>(%arg2)[%workgroup_size_x] + %16 = flow.dispatch.tensor.load %1, offsets = [0, 0, %arg2], sizes = [3, 3, %15], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:3x3x144xf32> -> tensor<3x3x?xf32> + %17 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 28)>(%arg0)[%workgroup_size_z] + %18 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 28)>(%arg1)[%workgroup_size_y] + %19 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 144)>(%arg2)[%workgroup_size_x] + %20 = affine.min affine_map<(d0)[s0] -> (-d0 + 28, s0)>(%arg0)[%workgroup_size_z] + %21 = affine.min affine_map<(d0)[s0] -> (-d0 + 28, s0)>(%arg1)[%workgroup_size_y] + %22 = affine.min affine_map<(d0)[s0] -> (-d0 + 144, s0)>(%arg2)[%workgroup_size_x] + %23 = linalg.init_tensor [1, %20, %21, %22] : tensor<1x?x?x?xf32> + %24 = linalg.fill(%cst, %23) : f32, tensor<1x?x?x?xf32> -> tensor<1x?x?x?xf32> + %25 = linalg.depthwise_conv2D_nhw {__internal_linalg_transform__ = "workgroup", dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%14, %16 : tensor<1x?x?x?xf32>, tensor<3x3x?xf32>) outs(%24 : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32> + flow.dispatch.tensor.store %25, %2, offsets = [0, %arg0, %arg1, %arg2], sizes = [1, %17, %18, %19], strides = [1, 1, 1, 1] : tensor<1x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:1x28x28x144xf32> + } + } + } + return + } + hal.interface private @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @dwconv_28x28x144 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [16, 8, 8]} +// CHECK-SAME: workgroup_size = [4 : index, 4 : index, 4 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[C9:.+]] = constant 9 : index +// CHECK-NEXT: %[[C3:.+]] = constant 3 : index +// CHECK-NEXT: hal.return %[[C9]], %[[C3]], %[[C3]] + +// CHECK: func @dwconv_28x28x144() +// CHECK: linalg.depthwise_conv2D_nhw +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[0, 8, 8, 16], [], [0, 2, 2, 4]]} + +// ----- + +// Depthwise conv - tiny OC/OW/OH - starving the GPU. + +hal.executable @dwconv_4x4x8 { + hal.interface public @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant public @vulkan_spirv_fb, target = #hal.executable.target<"vulkan", "vulkan-spirv-fb", { + spv.target_env = #spv.target_env<#spv.vce<v1.4, [Shader], []>, Qualcomm:IntegratedGPU, { + max_compute_shared_memory_size = 32768 : i32, + max_compute_workgroup_invocations = 1024 : i32, + max_compute_workgroup_size = dense<[1024, 1024, 64]> : vector<3xi32>, + subgroup_size = 64 : i32}> + }> { + hal.executable.entry_point public @dwconv_4x4x8 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @dwconv_4x4x8() { + %c0 = constant 0 : index + %c8 = constant 8 : index + %c4 = constant 4 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1x9x9x8xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:3x3x8xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1x4x4x8xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : index + %workgroup_size_z = hal.interface.workgroup.size[2] : index + %workgroup_id_x = hal.interface.workgroup.id[0] : index + %workgroup_count_x = hal.interface.workgroup.count[0] : index + %workgroup_id_y = hal.interface.workgroup.id[1] : index + %workgroup_count_y = hal.interface.workgroup.count[1] : index + %workgroup_id_z = hal.interface.workgroup.id[2] : index + %workgroup_count_z = hal.interface.workgroup.count[2] : index + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_z, %workgroup_size_z] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_z, %workgroup_size_z] + scf.for %arg0 = %3 to %c4 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg1 = %5 to %c4 step %6 { + %7 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %8 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg2 = %7 to %c8 step %8 { + %9 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg0) + %10 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 9)>(%arg0)[%workgroup_size_z] + %11 = affine.apply affine_map<(d0) -> (d0 * 2)>(%arg1) + %12 = affine.min affine_map<(d0)[s0] -> (s0 * 2 + 1, d0 * -2 + 9)>(%arg1)[%workgroup_size_y] + %13 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 8)>(%arg2)[%workgroup_size_x] + %14 = flow.dispatch.tensor.load %0, offsets = [0, %9, %11, %arg2], sizes = [1, %10, %12, %13], strides = [1, 1, 1, 1] : !flow.dispatch.tensor<readonly:1x9x9x8xf32> -> tensor<1x?x?x?xf32> + %15 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 8)>(%arg2)[%workgroup_size_x] + %16 = flow.dispatch.tensor.load %1, offsets = [0, 0, %arg2], sizes = [3, 3, %15], strides = [1, 1, 1] : !flow.dispatch.tensor<readonly:3x3x8xf32> -> tensor<3x3x?xf32> + %17 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 4)>(%arg0)[%workgroup_size_z] + %18 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 4)>(%arg1)[%workgroup_size_y] + %19 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 8)>(%arg2)[%workgroup_size_x] + %20 = affine.min affine_map<(d0)[s0] -> (-d0 + 4, s0)>(%arg0)[%workgroup_size_z] + %21 = affine.min affine_map<(d0)[s0] -> (-d0 + 4, s0)>(%arg1)[%workgroup_size_y] + %22 = affine.min affine_map<(d0)[s0] -> (-d0 + 8, s0)>(%arg2)[%workgroup_size_x] + %23 = linalg.init_tensor [1, %20, %21, %22] : tensor<1x?x?x?xf32> + %24 = linalg.fill(%cst, %23) : f32, tensor<1x?x?x?xf32> -> tensor<1x?x?x?xf32> + %25 = linalg.depthwise_conv2D_nhw {__internal_linalg_transform__ = "workgroup", dilations = dense<1> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>} ins(%14, %16 : tensor<1x?x?x?xf32>, tensor<3x3x?xf32>) outs(%24 : tensor<1x?x?x?xf32>) -> tensor<1x?x?x?xf32> + flow.dispatch.tensor.store %25, %2, offsets = [0, %arg0, %arg1, %arg2], sizes = [1, %17, %18, %19], strides = [1, 1, 1, 1] : tensor<1x?x?x?xf32> -> !flow.dispatch.tensor<writeonly:1x4x4x8xf32> + } + } + } + return + } + hal.interface private @io { + hal.interface.binding public @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding public @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @dwconv_4x4x8 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [8, 4, 4]} +// CHECK-SAME: workgroup_size = [2 : index, 4 : index, 4 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[C1:.+]] = constant 1 : index +// CHECK-NEXT: hal.return %[[C1]], %[[C1]], %[[C1]] + +// CHECK: func @dwconv_4x4x8() +// CHECK: linalg.depthwise_conv2D_nhw +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[0, 4, 4, 8], [], [0, 1, 1, 4]]}
diff --git a/iree/compiler/Codegen/SPIRV/test/tile_and_vectorize.mlir b/iree/compiler/Codegen/SPIRV/test/tile_and_vectorize.mlir index f0a227c..f1eb49b 100644 --- a/iree/compiler/Codegen/SPIRV/test/tile_and_vectorize.mlir +++ b/iree/compiler/Codegen/SPIRV/test/tile_and_vectorize.mlir
@@ -311,7 +311,7 @@ // CHECK: scf.for %[[IV5:.+]] = %[[TIDX]] to %{{.*}} step %[[BDIMX]] // CHECK: %[[OUT:.+]] = memref.subview %[[SV2]][0, %[[IV3]], %[[IV4]], %[[IV5]]] // CHECK: linalg.conv_2d_nhwc_hwcf -// CHECK-SAME: __internal_linalg_transform__ = "tile_conv_filter" +// CHECK-SAME: __internal_linalg_transform__ = "vectorize" // CHECK-SAME: outs(%[[OUT]] // -----