Merge pull request #6998 from ScottTodd:main-to-google PiperOrigin-RevId: 395576159
diff --git a/.github/workflows/build_package.yml b/.github/workflows/build_package.yml index 44103a9..79b7e9d 100644 --- a/.github/workflows/build_package.yml +++ b/.github/workflows/build_package.yml
@@ -263,4 +263,4 @@ workflow: Validate and Publish Release token: ${{ secrets.WRITE_ACCESS_TOKEN }} ref: "${{ env.tag_name }}" - inputs: '{"release_id": "${{ github.event.inputs.release_id }}"}' + inputs: '{"release_id": "${{ github.event.inputs.release_id }}", "package_version": "${{ github.event.inputs.package_version }}", "build_run_id": "${{ github.run_id }}"}'
diff --git a/.github/workflows/validate_and_publish_release.yml b/.github/workflows/validate_and_publish_release.yml index 2124c8d..a06b660 100644 --- a/.github/workflows/validate_and_publish_release.yml +++ b/.github/workflows/validate_and_publish_release.yml
@@ -1,5 +1,5 @@ -# Publishes a draft snapshot, after performing validation to make sure all -# packages are stable (coming soon). +# Publishes a draft snapshot, after performing validation to make sure all linux +# packages are stable. name: Validate and Publish Release on: @@ -8,13 +8,71 @@ release_id: description: 'Release id to publish' required: true - default: '' - -# TODO(jennik): Add more validation. https://github.com/google/iree/issues/6316 + package_version: + description: 'Version of the package' + required: true + build_run_id: + description: 'Run ID for the build_package.yml workflow that triggered this workflow' + required: true jobs: + validate_packages: + name: "Validate packages" + # TODO(jennik): Look into testing windows and macos builds. + runs-on: ubuntu-18.04 + steps: + - name: Download packages + id: download_packages + uses: dawidd6/action-download-artifact@v2 + with: + github_token: ${{secrets.WRITE_ACCESS_TOKEN}} + workflow: build_package.yml + run_id: ${{ github.event.inputs.build_run_id }} + - name: Extract and display downloaded files + run: | + tar -xf artifact/iree-dist-snapshot-${{ github.event.inputs.package_version }}-linux-x86_64.tar.xz + pwd + ls -R + - name: Set up python + id: set_up_python + uses: actions/setup-python@v2 + with: + python-version: '3.x' + - name: Install python packages + id: install_python_packages + run: | + python -m pip install -f file://$PWD/artifact/ iree-compiler-snapshot iree-jax-snapshot iree-runtime-snapshot iree-tools-tflite-snapshot iree-tools-tf-snapshot iree-tools-xla-snapshot + - name: Run iree-benchmark-module + id: run_iree_benchmark_module + run: ./bin/iree-benchmark-module --help + - name: Run iree-benchmark-trace + id: run_iree_benchmark_trace + run: ./bin/iree-benchmark-trace --help + - name: Run iree-flatcc-cli + id: run_iree_flatcc_cli + run: ./bin/iree-flatcc-cli --help + - name: Run iree-opt + id: run_iree_opt + run: ./bin/iree-opt --help + - name: Run iree-run-mlir + id: run_iree_run_mlir + run: ./bin/iree-run-mlir --help + - name: Run iree-run-module + id: run_iree_run_module + run: ./bin/iree-run-module --help + - name: Run iree-run-trace + id: run_iree_run_trace + run: ./bin/iree-run-trace --help + - name: Run iree-tblgen + id: run_iree_tblgen + run: ./bin/iree-tblgen --help + - name: Run iree-translate + id: run_iree_translate + run: ./bin/iree-translate --help + publish_release: name: "Publish release" + needs: validate_packages runs-on: ubuntu-18.04 steps: - name: Publish Release
diff --git a/bindings/python/iree/compiler/core.py b/bindings/python/iree/compiler/core.py index a74466a..d9baedc 100644 --- a/bindings/python/iree/compiler/core.py +++ b/bindings/python/iree/compiler/core.py
@@ -172,9 +172,6 @@ cl = [ iree_translate, input_file, - # People are only looking at stderr here if the tool fails and there's - # already lots of output, so adding this doesn't cost much. - "--print-main-address", f"--iree-input-type={options.input_type.value}", f"--iree-vm-bytecode-module-output-format={options.output_format.value}", ]
diff --git a/docs/developers/debugging/releases.md b/docs/developers/debugging/releases.md index efd7ba9..6f82ac3 100644 --- a/docs/developers/debugging/releases.md +++ b/docs/developers/debugging/releases.md
@@ -84,3 +84,50 @@ Release mode `main-dist` package above once, which will drop binaries in the `iree-install` directory. Then build the `py-runtime-pkg` or equiv and iterate further in the build directory. Ditto for TF/XLA/etc. + +## Testing releases on your fork + +To avoid interrupting the regular releases published on the IREE github, you +can test any changes to the release process on your own fork. Some setup is +required before these github actions will work on your fork and development +branch. + +To run +[`schedule_snapshot_release.yml`](https://github.com/google/iree/blob/main/.github/workflows/schedule_snapshot_release.yml), +comment out +[this line](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/schedule_snapshot_release.yml#L14): +```yaml +# Don't run this in everyone's forks. +if: github.repository == 'google/iree' +``` + +And change the branch from 'main' to the branch you are developing on +[here](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/schedule_snapshot_release.yml#L37): +```yaml +- name: Pushing changes + uses: ad-m/github-push-action@v0.6.0 + with: + github_token: ${{ secrets.WRITE_ACCESS_TOKEN }} + branch: main + tags: true +``` + +To speed up +[`build_package.yml`](https://github.com/google/iree/blob/main/.github/workflows/build_package.yml), +you may want to comment out some of the builds +[here](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L34-L87). +The +[`py-pure-pkgs`](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L52) +build takes only ~2 minutes and the +[`py-runtime-pkg`](https://github.com/google/iree/blob/392449e986493bf710e3da637ebf807715da9ffe/.github/workflows/build_package.yml#L39) +build takes ~5, while the others can take several hours. + +From your development branch, you can manually run the +[Schedule Snapshot Release](https://github.com/google/iree/actions/workflows/schedule_snapshot_release.yml) +action, which invokes the +[Build Native Release Packages](https://github.com/google/iree/actions/workflows/build_package.yml) +action, which finally invokes the +[Validate and Publish Release](https://github.com/google/iree/actions/workflows/validate_and_publish_release.yml) +action. If you already have a draft release and know the release id, package +version, and run ID from a previous Build Native Release Packages run, you can +also manually run just the Validate and Publish Release action.
diff --git a/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp b/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp index c12fac9..27deff8 100644 --- a/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp +++ b/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp
@@ -176,6 +176,15 @@ IREE::HAL::InterfaceBindingOp binding = input.queryBindingOp(); llvmInputTypes[argMapping[binding]] = llvmType; }); + // As a convention with HAL, push constants are appended as kernel arguments + // after all the binding inputs. + uint64_t numConstants = 0; + funcOp.walk([&](IREE::HAL::InterfaceLoadConstantOp constant) { + numConstants = + std::max(constant.offset().getZExtValue() + 1, numConstants); + }); + llvmInputTypes.resize(argMapping.size() + numConstants, + rewriter.getI32Type()); signatureConverter.addInputs(llvmInputTypes); // Construct newFunc with all attributes except return type & symbol name. @@ -256,14 +265,72 @@ rewriter, loc, *getTypeConverter(), memrefType, llvmBufferBasePtr); rewriter.replaceOp(op, {desc}); } else { - // TODO: pull those paramters from HAL constants. - assert(0 && "TODO: implement dynamic shape"); + ValueRange dynamicDims = adaptor.dynamic_dims(); + assert(memrefType.getNumDynamicDims() == dynamicDims.size()); + int64_t rank = memrefType.getRank(); + + // Build MemRef descriptor for this interface binding. + auto desc = MemRefDescriptor::undef( + rewriter, loc, typeConverter->convertType(memrefType)); + desc.setAllocatedPtr(rewriter, loc, llvmBufferBasePtr); + desc.setAlignedPtr(rewriter, loc, llvmBufferBasePtr); + desc.setConstantOffset(rewriter, loc, 0); + + // Update memref descriptor shape. Dynamic dimensions can be mixed with + // static dimensions, like [128, ?, 128]. + int dynamicDimIndex = 0; + for (int i = 0; i < rank; ++i) { + if (memrefType.isDynamicDim(i)) { + desc.setSize(rewriter, loc, i, dynamicDims[dynamicDimIndex++]); + } else { + desc.setConstantSize(rewriter, loc, i, memrefType.getDimSize(i)); + } + } + + // Compute and update strides. Assume that MemRefs are row-major, that is, + // following index linearization: + // x[i, j, k] = i * x.dim[1] * x.dim[2] + j * x.dim[2] + k + desc.setConstantStride(rewriter, loc, rank - 1, 1); + for (int i = rank - 2; i >= 0; --i) { + auto stride = desc.stride(rewriter, loc, i + 1); + auto dim = desc.size(rewriter, loc, i + 1); + Value strideVal = rewriter.create<LLVM::MulOp>(loc, stride, dim); + desc.setStride(rewriter, loc, i, strideVal); + } + rewriter.replaceOp(op, {desc}); } return success(); } }; +class ConvertIREEConstantOp : public ConvertToLLVMPattern { + public: + explicit ConvertIREEConstantOp(MLIRContext *context, + LLVMTypeConverter &converter) + : ConvertToLLVMPattern( + IREE::HAL::InterfaceLoadConstantOp::getOperationName(), context, + converter) {} + LogicalResult matchAndRewrite( + Operation *op, ArrayRef<Value> operands, + ConversionPatternRewriter &rewriter) const override { + // Bail until nested under an LLVMFuncOp. + auto llvmFuncOp = op->getParentOfType<LLVM::LLVMFuncOp>(); + if (!llvmFuncOp) return failure(); + assert(llvmFuncOp.getNumArguments() > 0); + + llvm::SmallDenseMap<Operation *, size_t> argMapping = + getKernelArgMapping(llvmFuncOp); + auto ireeConstantOp = cast<IREE::HAL::InterfaceLoadConstantOp>(op); + mlir::BlockArgument llvmBufferArg = llvmFuncOp.getArgument( + argMapping.size() + ireeConstantOp.offset().getZExtValue()); + assert(llvmBufferArg.getType().isInteger(32)); + Type dstType = getTypeConverter()->convertType(ireeConstantOp.getType()); + rewriter.replaceOpWithNewOp<LLVM::ZExtOp>(op, dstType, llvmBufferArg); + return success(); + } +}; + /// A pattern to convert hal.interface.workgroup.id/count/size into /// corresponding NVVM/ROCDL ops. template <typename InterfaceOpTy, typename XOp, typename YOp, typename ZOp> @@ -305,7 +372,8 @@ OwningRewritePatternList &patterns, LLVMTypeConverter &converter, bool useROCM) { - patterns.insert<ConvertFunc, ConvertIREEBindingOp>(context, converter); + patterns.insert<ConvertFunc, ConvertIREEBindingOp, ConvertIREEConstantOp>( + context, converter); if (useROCM) { patterns.insert<HALInterfaceWorkgroupOpsConverter< IREE::HAL::InterfaceWorkgroupIDOp, ROCDL::BlockIdXOp,
diff --git a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp index a069bed..bcb7bf5 100644 --- a/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp +++ b/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -43,8 +43,8 @@ // Infer the MxN size of the matmul based on operands and indexing maps. auto lhsShape = getUntiledShape(op.getInputOperand(0)->get()); auto rhsShape = getUntiledShape(op.getInputOperand(1)->get()); - int64_t sizeM = -1; - int64_t sizeN = -1; + int64_t sizeM = ShapedType::kDynamicSize; + int64_t sizeN = ShapedType::kDynamicSize; auto outputMap = op.getTiedIndexingMap(op.getOutputOperand(0)); for (unsigned i = 0; i < lhsShape.size(); i++) { if (op.getTiedIndexingMap(op.getInputOperand(0)).getDimPosition(i) == @@ -65,25 +65,29 @@ 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); - // Pick the best configuration where the original shape is aligned on the tile - // size. - for (TileWorkgroupSizePair &config : tileSizeConfig) { - if (sizeN % config.tileSize[1] == 0 && sizeM % config.tileSize[0] == 0) { - tileX = config.tileSize[0]; - tileY = config.tileSize[1]; - tileK = config.tileSize[2]; - workgroupSize.assign(config.workgroupSize.begin(), - config.workgroupSize.end()); - break; + bool isStaticSize = + sizeM != ShapedType::kDynamicSize && sizeN != ShapedType::kDynamicSize; + if (isStaticSize) { + // 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); + // Pick the best configuration where the original shape is aligned on the + // tile size. + for (TileWorkgroupSizePair &config : tileSizeConfig) { + if (sizeN % config.tileSize[1] == 0 && sizeM % config.tileSize[0] == 0) { + tileX = config.tileSize[0]; + tileY = config.tileSize[1]; + tileK = config.tileSize[2]; + workgroupSize.assign(config.workgroupSize.begin(), + config.workgroupSize.end()); + break; + } } } // Currently just a basic tile size to enable tiling and vectorization.
diff --git a/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir b/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir index 00e5c31..7ee0fcc 100644 --- a/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir +++ b/iree/compiler/Codegen/LLVMGPU/test/convert_to_nvvm.mlir
@@ -1,4 +1,4 @@ -// RUN: iree-opt -iree-convert-to-nvvm %s | IreeFileCheck %s +// RUN: iree-opt -iree-convert-to-nvvm -split-input-file %s | IreeFileCheck %s // Test that that standard and GPU ops are converted to LLVM and NVVM. func @abs_ex_dispatch_0() { @@ -37,3 +37,44 @@ // CHECK: llvm.insertvalue %[[PTR]], %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)> // CHECK: nvvm.read.ptx.sreg.tid.x // CHECK: llvm.fadd + +// ----- + +func @abs_dynamic() { + %c0 = constant 0 : index + %c128 = constant 128 : index + %s = hal.interface.load.constant offset = 1 : index + %0 = hal.interface.binding.subspan @io::@arg0[%c128] : memref<?xf32>{%s} + %1 = hal.interface.binding.subspan @io::@arg1[%c0] : memref<16xi32> + %2 = hal.interface.binding.subspan @io::@ret0[%c0] : memref<16xf32> + %3 = "gpu.block_id"() {dimension = "x"} : () -> index + %4 = "gpu.block_dim"() {dimension = "x"} : () -> index + %5 = "gpu.thread_id"() {dimension = "x"} : () -> index + %6 = muli %3, %4 : index + %7 = addi %6, %5 : index + %9 = memref.load %0[%7] : memref<?xf32> + %10 = memref.load %1[%7] : memref<16xi32> + %11 = sitofp %10 : i32 to f32 + %12 = addf %9, %11 : f32 + memref.store %12, %2[%7] : memref<16xf32> + return +} +hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @arg0, set=0, binding=4, type="StorageBuffer", access="Read" + hal.interface.binding @arg1, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @ret0, set=1, binding=2, type="StorageBuffer", access="Write|Discard" +} + +// CHECK-LABEL: llvm.func @abs_dynamic +// CHECK-SAME: (%[[ARG0:.+]]: !llvm.ptr<i32> {llvm.align = 16 : i32}, +// CHECK-SAME: %[[ARG1:.+]]: !llvm.ptr<f32> {llvm.align = 16 : i32}, %[[ARG2:.+]]: !llvm.ptr<f32> {llvm.align = 16 : i32}, +// CHECK-SAME: %[[ARG3:.+]]: i32, %[[ARG4:.+]]: i32) +// CHECK: %[[C128:.+]] = llvm.mlir.constant(128 : index) : i64 +// CHECK: %{{.*}} = llvm.zext %[[ARG4]] : i32 to i64 +// CHECK: %[[PTRI8:.+]] = llvm.bitcast %[[ARG1]] : !llvm.ptr<f32> to !llvm.ptr<i8> +// CHECK: %[[OFF:.+]] = llvm.getelementptr %[[PTRI8]][%[[C128]]] : (!llvm.ptr<i8>, i64) -> !llvm.ptr<i8> +// CHECK: %[[PTR:.+]] = llvm.bitcast %[[OFF]] : !llvm.ptr<i8> to !llvm.ptr<f32> +// CHECK: llvm.insertvalue %[[PTR]], %{{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)> +// CHECK: llvm.insertvalue %[[PTR]], %{{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>)> +// CHECK: nvvm.read.ptx.sreg.tid.x +// CHECK: llvm.fadd
diff --git a/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp new file mode 100644 index 0000000..c8b441f --- /dev/null +++ b/iree/compiler/Codegen/SPIRV/AdrenoConfig.cpp
@@ -0,0 +1,125 @@ +// Copyright 2021 The IREE Authors +// +// Licensed under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//===- AdrenoConfig.h - Adreno CodeGen Configurations ---------------------===// +// +// This file contains CodeGen configurations for Adreno GPUs. +// +//===----------------------------------------------------------------------===// + +#include <array> + +#include "iree/compiler/Codegen/SPIRV/KernelConfig.h" +#include "iree/compiler/Codegen/Transforms/Transforms.h" +#include "iree/compiler/Codegen/Utils/Utils.h" +#include "llvm/Support/MathExtras.h" +#include "mlir/Dialect/Linalg/IR/LinalgOps.h" + +namespace mlir { +namespace iree_compiler { +namespace detail { + +//===----------------------------------------------------------------------===// +// Matmul +//===----------------------------------------------------------------------===// + +static LogicalResult setOpConfig(linalg::MatmulOp op) { + ArrayRef<int64_t> lhsShape = getUntiledShape(op.inputs()[0]); + ArrayRef<int64_t> rhsShape = getUntiledShape(op.inputs()[1]); + if (llvm::any_of(lhsShape, ShapedType::isDynamic)) return success(); + if (llvm::any_of(rhsShape, ShapedType::isDynamic)) return success(); + + int64_t dimM = lhsShape[0], dimK = lhsShape[1], dimN = rhsShape[1]; + + // The core idea is to distribute the matmul M/N dimension to the workgroup + // Y/X dimension, with each thread in a workgroup handling multiple vector + // elements. We start from the best (X, Y) and the tiling sizes for (M, N, K) + // and try different configurations by scaling them down until we find a + // configuration that can perfectly tile the input matmul. + + const int64_t bestX = 32, bestY = 2; + const int64_t bestThreadM = 8, bestThreadN = 8, bestThreadK = 4; + + int64_t residualThreads = bestX * bestY; + int64_t residualTilingFactor = (bestThreadM + bestThreadK) * bestThreadN; + + SmallVector<int64_t, 3> workgroupSize(3, 1); // (X, Y, Z) + SmallVector<int64_t, 4> workgroupTileSizes(3, 0); // (M, N, K) + SmallVector<int64_t, 4> invocationTileSizes(3, 0); // (M, N, K) + + // 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. + int64_t chosenTileSize = 4; + if (dimN % (x * chosenTileSize) == 0) { + workgroupSize[0] = x; + workgroupTileSizes[1] = x * chosenTileSize; + invocationTileSizes[1] = chosenTileSize; + residualThreads /= x; + assert(residualTilingFactor % chosenTileSize == 0); + residualTilingFactor /= chosenTileSize; + break; + } + } + if (workgroupTileSizes[1] == 0) return success(); + + // Deduce the configuration for the M dimension. Start with the best workgroup + // Y size, and reduce by a factor of two each time. + for (int64_t y = residualThreads; y >= 1; y >>= 1) { + int64_t chosenTileSize = 0; + // Reduce the thread tiling size by one each time. We read one row each + // time; so it's fine to not be some power of two here. + for (int64_t t = bestThreadM; t >= 1; --t) { + if (dimM % (y * t) == 0) { + chosenTileSize = t; + break; + } + } + if (chosenTileSize) { + workgroupSize[1] = y; + workgroupTileSizes[0] = y * chosenTileSize; + invocationTileSizes[0] = chosenTileSize; + assert(residualTilingFactor > chosenTileSize); + residualTilingFactor -= chosenTileSize; + break; + } + } + if (workgroupTileSizes[0] == 0) return success(); + + // Deduce the configuration for the K dimension. We need some power of two + // here so that we can do vector load. + for (int64_t t = llvm::PowerOf2Floor(residualTilingFactor); t >= 1; t >>= 1) { + if (dimK % t == 0) { + workgroupTileSizes[2] = invocationTileSizes[2] = t; + break; + } + } + + auto pipeline = IREE::HAL::DispatchLoweringPassPipeline::SPIRVVectorize; + TileSizesListType tileSizes; + tileSizes.push_back(workgroupTileSizes); + tileSizes.emplace_back(); + tileSizes.push_back(invocationTileSizes); + return setOpConfigAndEntryPointFnTranslation(op->getParentOfType<FuncOp>(), + op, tileSizes, {}, pipeline, + workgroupSize); +} + +//===----------------------------------------------------------------------===// +// Entry Point +//===----------------------------------------------------------------------===// + +LogicalResult setAdrenoCodeGenConfig(const spirv::TargetEnv &, + Operation *rootOp) { + return TypeSwitch<Operation *, LogicalResult>(rootOp) + .Case<linalg::MatmulOp>([](auto op) { return setOpConfig(op); }) + .Default([](Operation *) { return success(); }); +} + +} // namespace detail +} // namespace iree_compiler +} // namespace mlir
diff --git a/iree/compiler/Codegen/SPIRV/BUILD b/iree/compiler/Codegen/SPIRV/BUILD index 4d68549..4915dce 100644 --- a/iree/compiler/Codegen/SPIRV/BUILD +++ b/iree/compiler/Codegen/SPIRV/BUILD
@@ -13,6 +13,7 @@ cc_library( name = "SPIRV", srcs = [ + "AdrenoConfig.cpp", "ConvertToSPIRVPass.cpp", "KernelConfig.cpp", "MaliConfig.cpp",
diff --git a/iree/compiler/Codegen/SPIRV/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/CMakeLists.txt index 313e988..3cadce3 100644 --- a/iree/compiler/Codegen/SPIRV/CMakeLists.txt +++ b/iree/compiler/Codegen/SPIRV/CMakeLists.txt
@@ -18,6 +18,7 @@ "MemorySpace.h" "Utils.h" SRCS + "AdrenoConfig.cpp" "ConvertToSPIRVPass.cpp" "KernelConfig.cpp" "MaliConfig.cpp"
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp index a2c536a..d7e72f3 100644 --- a/iree/compiler/Codegen/SPIRV/KernelConfig.cpp +++ b/iree/compiler/Codegen/SPIRV/KernelConfig.cpp
@@ -252,6 +252,9 @@ case spirv::Vendor::NVIDIA: result = detail::setNVIDIACodeGenConfig(targetEnv, rootOp); break; + case spirv::Vendor::Qualcomm: + result = detail::setAdrenoCodeGenConfig(targetEnv, rootOp); + break; default: break; }
diff --git a/iree/compiler/Codegen/SPIRV/KernelConfig.h b/iree/compiler/Codegen/SPIRV/KernelConfig.h index 68660c0..61f2a13 100644 --- a/iree/compiler/Codegen/SPIRV/KernelConfig.h +++ b/iree/compiler/Codegen/SPIRV/KernelConfig.h
@@ -26,21 +26,20 @@ namespace detail { -/// If the given `rootOp` has known good CodeGen configuration for Mali GPUs, -/// attaches a `translation.info` attribute to the entry point containing -/// `rootOp` and a `lowering.config` attribute to `rootOp`. +/// Sets CodeGen configuration for GPUs from a specific vendor. +/// +/// If the given `rootOp` has known good CodeGen configuration, attaches a +/// `translation.info` attribute to the entry point containing `rootOp` and a +/// `lowering.config` attribute to `rootOp`. +/// /// Returns success when either no configuration is found or a configuration is /// successfullly attached as attribute. Returns failure only when there is an /// issue attaching the attribute. + +LogicalResult setAdrenoCodeGenConfig(const spirv::TargetEnv &targetEnv, + Operation *rootOp); LogicalResult setMaliCodeGenConfig(const spirv::TargetEnv &targetEnv, Operation *rootOp); - -/// If the given `rootOp` has known good CodeGen configuration for NVIDIA GPUs, -/// attaches a `translation.info` attribute to the entry point containing -/// `rootOp` and a `lowering.config` attribute to `rootOp`. -/// Returns success when either no configuration is found or a configuration is -/// successfullly attached as attribute. Returns failure only when there is an -/// issue attaching the attribute. LogicalResult setNVIDIACodeGenConfig(const spirv::TargetEnv &targetEnv, Operation *rootOp);
diff --git a/iree/compiler/Codegen/SPIRV/test/BUILD b/iree/compiler/Codegen/SPIRV/test/BUILD index fc4402b..330e6a3 100644 --- a/iree/compiler/Codegen/SPIRV/test/BUILD +++ b/iree/compiler/Codegen/SPIRV/test/BUILD
@@ -19,13 +19,14 @@ name = "lit", srcs = enforce_glob( [ + "config_adreno_matmul.mlir", + "config_linalg_ext_ops.mlir", "convert_to_spirv.mlir", "distribute_to_global_id.mlir", "fold_gpu_procid_uses.mlir", "pipeline_matmul_cooperative_matrix.mlir", "pipeline_matmul_vectorization.mlir", "remove_one_trip_tiled_loop.mlir", - "set_lowering_config.mlir", "tile_and_distribute_scatter.mlir", "tile_and_distribute_sort.mlir", "tile_and_vectorize.mlir",
diff --git a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt index 31ccf04..ef888fb 100644 --- a/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt +++ b/iree/compiler/Codegen/SPIRV/test/CMakeLists.txt
@@ -14,13 +14,14 @@ NAME lit SRCS + "config_adreno_matmul.mlir" + "config_linalg_ext_ops.mlir" "convert_to_spirv.mlir" "distribute_to_global_id.mlir" "fold_gpu_procid_uses.mlir" "pipeline_matmul_cooperative_matrix.mlir" "pipeline_matmul_vectorization.mlir" "remove_one_trip_tiled_loop.mlir" - "set_lowering_config.mlir" "tile_and_distribute_scatter.mlir" "tile_and_distribute_sort.mlir" "tile_and_vectorize.mlir"
diff --git a/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir b/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir new file mode 100644 index 0000000..aee34bd --- /dev/null +++ b/iree/compiler/Codegen/SPIRV/test/config_adreno_matmul.mlir
@@ -0,0 +1,379 @@ +// 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 + +// Large matmul that can match the best tiling scheme. + +hal.executable @matmul_1024x2048x512 { + hal.interface @io { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant @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 @matmul_1024x2048x512 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @matmul_1024x2048x512() { + %c0 = constant 0 : index + %c2048 = constant 2048 : index + %c1024 = constant 1024 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:1024x512xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:512x2048xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:1024x2048xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : 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 + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg0 = %3 to %c1024 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg1 = %5 to %c2048 step %6 { + %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 1024)>(%arg0)[%workgroup_size_y] + %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 512], strides = [1, 1] : !flow.dispatch.tensor<readonly:1024x512xf32> -> tensor<?x512xf32> + %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 2048)>(%arg1)[%workgroup_size_x] + %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [512, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:512x2048xf32> -> tensor<512x?xf32> + %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 1024)>(%arg0)[%workgroup_size_y] + %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 2048)>(%arg1)[%workgroup_size_x] + %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 1024, s0)>(%arg0)[%workgroup_size_y] + %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 2048, s0)>(%arg1)[%workgroup_size_x] + %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32> + %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32> + %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x512xf32>, tensor<512x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32> + flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:1024x2048xf32> + } + } + return + } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @matmul_1024x2048x512 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [128, 16]} +// CHECK-SAME: workgroup_size = [32 : index, 2 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[ONE:.+]] = constant 1 : index +// CHECK-NEXT: %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 128)>()[%[[X]]] +// CHECK-NEXT: %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 16)>()[%[[Y]]] +// CHECK-NEXT: hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]] + +// CHECK: func @matmul_1024x2048x512() +// CHECK: linalg.matmul +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[16, 128, 16], [], [8, 4, 16]]} + +// ----- + +// Small matmul N that can still tile to all threads in a workgroup. + +hal.executable @matmul_3136x24x96 { + hal.interface @io { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant @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 @matmul_3136x24x96 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @matmul_3136x24x96() { + %c0 = constant 0 : index + %c24 = constant 24 : index + %c3136 = constant 3136 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:3136x96xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:96x24xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:3136x24xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : 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 + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg0 = %3 to %c3136 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg1 = %5 to %c24 step %6 { + %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 3136)>(%arg0)[%workgroup_size_y] + %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 96], strides = [1, 1] : !flow.dispatch.tensor<readonly:3136x96xf32> -> tensor<?x96xf32> + %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 24)>(%arg1)[%workgroup_size_x] + %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [96, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:96x24xf32> -> tensor<96x?xf32> + %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 3136)>(%arg0)[%workgroup_size_y] + %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 24)>(%arg1)[%workgroup_size_x] + %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 3136, s0)>(%arg0)[%workgroup_size_y] + %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 24, s0)>(%arg1)[%workgroup_size_x] + %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32> + %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32> + %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x96xf32>, tensor<96x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32> + flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:3136x24xf32> + } + } + return + } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @matmul_3136x24x96 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [8, 224]} +// CHECK-SAME: workgroup_size = [2 : index, 32 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[ONE:.+]] = constant 1 : index +// CHECK-NEXT: %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 8)>()[%[[X]]] +// CHECK-NEXT: %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 224)>()[%[[Y]]] +// CHECK-NEXT: hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]] + +// CHECK: func @matmul_3136x24x96() +// CHECK: linalg.matmul +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[224, 8, 16], [], [7, 4, 16]]} + +// ----- + +// Small matmul M that can still tile to all threads in a workgroup. + +hal.executable @matmul_196x64x192 { + hal.interface @io { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant @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 @matmul_196x64x192 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @matmul_196x64x192() { + %c0 = constant 0 : index + %c64 = constant 64 : index + %c196 = constant 196 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:196x192xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:192x64xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:196x64xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : 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 + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg0 = %3 to %c196 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg1 = %5 to %c64 step %6 { + %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 196)>(%arg0)[%workgroup_size_y] + %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 192], strides = [1, 1] : !flow.dispatch.tensor<readonly:196x192xf32> -> tensor<?x192xf32> + %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 64)>(%arg1)[%workgroup_size_x] + %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [192, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:192x64xf32> -> tensor<192x?xf32> + %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 196)>(%arg0)[%workgroup_size_y] + %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 64)>(%arg1)[%workgroup_size_x] + %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 196, s0)>(%arg0)[%workgroup_size_y] + %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 64, s0)>(%arg1)[%workgroup_size_x] + %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32> + %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32> + %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x192xf32>, tensor<192x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32> + flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:196x64xf32> + } + } + return + } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @matmul_196x64x192 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [64, 28]} +// CHECK-SAME: workgroup_size = [16 : index, 4 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[ONE:.+]] = constant 1 : index +// CHECK-NEXT: %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 64)>()[%[[X]]] +// CHECK-NEXT: %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 28)>()[%[[Y]]] +// CHECK-NEXT: hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]] + +// CHECK: func @matmul_196x64x192() +// CHECK: linalg.matmul +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[28, 64, 16], [], [7, 4, 16]]} + +// ----- + +// Small matmul K that can still tile to all threads in a workgroup. + +hal.executable @matmul_12544x96x16 { + hal.interface @io { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant @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 @matmul_12544x96x16 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @matmul_12544x96x16() { + %c0 = constant 0 : index + %c96 = constant 96 : index + %c12544 = constant 12544 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : memref<12544x16xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : memref<16x96xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : memref<12544x96xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : 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 + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg0 = %3 to %c12544 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg1 = %5 to %c96 step %6 { + %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 12544)>(%arg0)[%workgroup_size_y] + %8 = memref.subview %0[%arg0, 0] [%7, 16] [1, 1] : memref<12544x16xf32> to memref<?x16xf32, affine_map<(d0, d1)[s0] -> (d0 * 16 + s0 + d1)>> + %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 96)>(%arg1)[%workgroup_size_x] + %10 = memref.subview %1[0, %arg1] [16, %9] [1, 1] : memref<16x96xf32> to memref<16x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>> + %11 = memref.subview %2[%arg0, %arg1] [%7, %9] [1, 1] : memref<12544x96xf32> to memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>> + linalg.fill(%cst, %11) : f32, memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>> + linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : memref<?x16xf32, affine_map<(d0, d1)[s0] -> (d0 * 16 + s0 + d1)>>, memref<16x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>) outs(%11 : memref<?x?xf32, affine_map<(d0, d1)[s0] -> (d0 * 96 + s0 + d1)>>) + } + } + return + } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @matmul_12544x96x16 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [32, 64]} +// CHECK-SAME: workgroup_size = [8 : index, 8 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[ONE:.+]] = constant 1 : index +// CHECK-NEXT: %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%[[X]]] +// CHECK-NEXT: %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 64)>()[%[[Y]]] +// CHECK-NEXT: hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]] + +// CHECK: func @matmul_12544x96x16() +// CHECK: linalg.matmul +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[64, 32, 16], [], [8, 4, 16]]} + +// ----- + +// Odd matmul M and small N that cannot utilize all threads in a workgroup. + +hal.executable @matmul_49x160x576 { + hal.interface @io { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + hal.executable.variant @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 @matmul_49x160x576 attributes {interface = @io, ordinal = 0 : index} + builtin.module { + func @matmul_49x160x576() { + %c0 = constant 0 : index + %c160 = constant 160 : index + %c49 = constant 49 : index + %cst = constant 0.000000e+00 : f32 + %0 = hal.interface.binding.subspan @io::@s0b0_ro_external[%c0] : !flow.dispatch.tensor<readonly:49x576xf32> + %1 = hal.interface.binding.subspan @io::@s0b1_ro_external[%c0] : !flow.dispatch.tensor<readonly:576x160xf32> + %2 = hal.interface.binding.subspan @io::@s0b2_xw_external[%c0] : !flow.dispatch.tensor<writeonly:49x160xf32> + %workgroup_size_x = hal.interface.workgroup.size[0] : index + %workgroup_size_y = hal.interface.workgroup.size[1] : 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 + %3 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_y, %workgroup_size_y] + %4 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_y, %workgroup_size_y] + scf.for %arg0 = %3 to %c49 step %4 { + %5 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_id_x, %workgroup_size_x] + %6 = affine.apply affine_map<()[s0, s1] -> (s0 * s1)>()[%workgroup_count_x, %workgroup_size_x] + scf.for %arg1 = %5 to %c160 step %6 { + %7 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 49)>(%arg0)[%workgroup_size_y] + %8 = flow.dispatch.tensor.load %0, offsets = [%arg0, 0], sizes = [%7, 576], strides = [1, 1] : !flow.dispatch.tensor<readonly:49x576xf32> -> tensor<?x576xf32> + %9 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 160)>(%arg1)[%workgroup_size_x] + %10 = flow.dispatch.tensor.load %1, offsets = [0, %arg1], sizes = [576, %9], strides = [1, 1] : !flow.dispatch.tensor<readonly:576x160xf32> -> tensor<576x?xf32> + %11 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 49)>(%arg0)[%workgroup_size_y] + %12 = affine.min affine_map<(d0)[s0] -> (s0, -d0 + 160)>(%arg1)[%workgroup_size_x] + %13 = affine.min affine_map<(d0)[s0] -> (-d0 + 49, s0)>(%arg0)[%workgroup_size_y] + %14 = affine.min affine_map<(d0)[s0] -> (-d0 + 160, s0)>(%arg1)[%workgroup_size_x] + %15 = linalg.init_tensor [%13, %14] : tensor<?x?xf32> + %16 = linalg.fill(%cst, %15) : f32, tensor<?x?xf32> -> tensor<?x?xf32> + %17 = linalg.matmul {__internal_linalg_transform__ = "workgroup"} ins(%8, %10 : tensor<?x576xf32>, tensor<576x?xf32>) outs(%16 : tensor<?x?xf32>) -> tensor<?x?xf32> + flow.dispatch.tensor.store %17, %2, offsets = [%arg0, %arg1], sizes = [%11, %12], strides = [1, 1] : tensor<?x?xf32> -> !flow.dispatch.tensor<writeonly:49x160xf32> + } + } + return + } + hal.interface @io attributes {sym_visibility = "private"} { + hal.interface.binding @s0b0_ro_external, set=0, binding=0, type="StorageBuffer", access="Read" + hal.interface.binding @s0b1_ro_external, set=0, binding=1, type="StorageBuffer", access="Read" + hal.interface.binding @s0b2_xw_external, set=0, binding=2, type="StorageBuffer", access="Write|Discard" + } + } + } +} + +// CHECK-LABEL: hal.executable.entry_point public @matmul_49x160x576 +// CHECK-SAME: translation.info = {passPipeline = 6 : i32, workloadPerWorkgroup = [32, 7]} +// CHECK-SAME: workgroup_size = [8 : index, 1 : index, 1 : index] +// CHECK-NEXT: ^{{.+}}(%[[X:.+]]: index, %[[Y:.+]]: index, %{{.+}}: index): +// CHECK-NEXT: %[[ONE:.+]] = constant 1 : index +// CHECK-NEXT: %[[X_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 32)>()[%[[X]]] +// CHECK-NEXT: %[[Y_COUNT:.+]] = affine.apply affine_map<()[s0] -> (s0 ceildiv 7)>()[%[[Y]]] +// CHECK-NEXT: hal.return %[[X_COUNT]], %[[Y_COUNT]], %[[ONE]] + +// CHECK: func @matmul_49x160x576() +// CHECK: linalg.matmul +// CHECK-SAME{LITERAL}: lowering.config = {tileSizes = [[7, 32, 16], [], [7, 4, 16]]}
diff --git a/iree/compiler/Codegen/SPIRV/test/set_lowering_config.mlir b/iree/compiler/Codegen/SPIRV/test/config_linalg_ext_ops.mlir similarity index 100% rename from iree/compiler/Codegen/SPIRV/test/set_lowering_config.mlir rename to iree/compiler/Codegen/SPIRV/test/config_linalg_ext_ops.mlir
diff --git a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp index 87affe9..975a41b 100644 --- a/iree/compiler/Dialect/Flow/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/Flow/Transforms/Passes.cpp
@@ -104,8 +104,7 @@ // The generated ABI wrappers assume such an expansion and will generate code // to produce it from the original reflection metadata captured in the // previous pass. - passManager.addNestedPass<mlir::FuncOp>( - Shape::createExpandFunctionDynamicDimsPass()); + passManager.addPass(Shape::createExpandFunctionDynamicDimsPass()); // Special case peephole optimizations. {
diff --git a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp index 3b3df2c..fc00b5d 100644 --- a/iree/compiler/Dialect/HAL/Transforms/Passes.cpp +++ b/iree/compiler/Dialect/HAL/Transforms/Passes.cpp
@@ -96,8 +96,7 @@ // on explicit shape types (such as ranked_shape). After this pass, these // composite types will be expanded to primitives (i.e. one 'index' for each // dynamic dim in the case of ranked_shape). - passManager.addNestedPass<FuncOp>( - Shape::createExpandFunctionRankedShapeDimsPass()); + passManager.addPass(Shape::createExpandFunctionRankedShapeDimsPass()); passManager.addNestedPass<FuncOp>(createCanonicalizerPass()); passManager.addNestedPass<FuncOp>(createCSEPass());
diff --git a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp index 316cd72..edd0814 100644 --- a/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp +++ b/iree/compiler/Dialect/Shape/Transforms/FunctionSignatureExpansionPass.cpp
@@ -21,10 +21,95 @@ namespace mlir { namespace iree_compiler { namespace Shape { + +static LogicalResult expandCallGraphTypesForCall( + CallOp callOp, const TypeExpander &typeExpander) { + Location loc = callOp.getLoc(); + + OpBuilder builder(callOp); + // Create the operands list of the new `CallOp`. + SmallVector<Value> newOperands; + if (failed(typeExpander.expandSourceValuesToTarget( + loc, llvm::to_vector<6>(callOp.getOperands()), newOperands, + builder))) { + return failure(); + } + + // Create the new result types for the new `CallOp` and track the indices + // in the new call op's results that correspond to the old call op's + // results. + // + // expandedResultIndices[i] = "list of new result indices that old result + // i expanded to". + SmallVector<Type> newResultTypes; + SmallVector<SmallVector<unsigned>> expandedResultIndices; + for (Type resultType : callOp.getResultTypes()) { + unsigned oldSize = newResultTypes.size(); + if (failed(typeExpander.convertType(resultType, newResultTypes))) { + return failure(); + } + auto &resultMapping = expandedResultIndices.emplace_back(); + for (unsigned i = oldSize, e = newResultTypes.size(); i < e; i++) + resultMapping.push_back(i); + } + + CallOp newCallOp = builder.create<CallOp>(loc, callOp.getCallee(), + newResultTypes, newOperands); + + // Build a replacement value for each result to replace its uses. If a + // result has multiple mapping values, it needs to be materialized as a + // single value. + SmallVector<Value, 2> replacedValues; + replacedValues.reserve(callOp.getNumResults()); + for (unsigned i = 0, e = callOp.getNumResults(); i < e; ++i) { + auto decomposedValues = llvm::to_vector<6>( + llvm::map_range(expandedResultIndices[i], + [&](unsigned i) { return newCallOp.getResult(i); })); + if (decomposedValues.empty()) { + // No replacement is required. + replacedValues.push_back(nullptr); + } else if (decomposedValues.size() == 1) { + replacedValues.push_back(decomposedValues.front()); + } else { + Value materialized = typeExpander.castToSource(loc, callOp.getType(i), + decomposedValues, builder); + replacedValues.push_back(materialized); + } + } + callOp.replaceAllUsesWith(replacedValues); + return success(); +} + +static LogicalResult expandCallGraphTypes(ModuleOp module, + const TypeExpander &typeExpander) { + for (auto funcOp : module.getOps<FuncOp>()) { + OpBuilder builder(funcOp); + if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) || + failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>( + funcOp, builder))) { + return funcOp.emitError() + << "couldn't expand function signature or return"; + } + bool hadError = false; + SmallVector<Operation *> toErase; + funcOp.walk([&](CallOp callOp) { + if (failed(expandCallGraphTypesForCall(callOp, typeExpander))) { + hadError = true; + return; + } + toErase.push_back(callOp); + }); + if (hadError) return failure(); + for (Operation *op : toErase) op->erase(); + } + return success(); +} + namespace { class ExpandFunctionDynamicDimsPass - : public PassWrapper<ExpandFunctionDynamicDimsPass, FunctionPass> { + : public PassWrapper<ExpandFunctionDynamicDimsPass, + OperationPass<ModuleOp>> { void getDependentDialects(DialectRegistry ®istry) const override { registry.insert<ShapeDialect>(); } @@ -37,20 +122,16 @@ return "Expands dynamic dimensions in function signatures."; } - void runOnFunction() override { - auto funcOp = getFunction(); + void runOnOperation() override { auto &typeExpander = getDynamicShapeTypeExpander(); - OpBuilder builder(funcOp); - if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) || - failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>( - funcOp, builder))) { - return signalPassFailure(); - } + if (failed(expandCallGraphTypes(getOperation(), typeExpander))) + signalPassFailure(); } }; class ExpandFunctionRankedShapeDimsPass - : public PassWrapper<ExpandFunctionRankedShapeDimsPass, FunctionPass> { + : public PassWrapper<ExpandFunctionRankedShapeDimsPass, + OperationPass<ModuleOp>> { void getDependentDialects(DialectRegistry ®istry) const override { registry.insert<ShapeDialect>(); } @@ -63,15 +144,10 @@ return "Expands ranked_shape types at function boundaries to loose dims."; } - void runOnFunction() override { - auto funcOp = getFunction(); + void runOnOperation() override { auto &typeExpander = getShapeToPrimitiveTypeExpander(); - OpBuilder builder(funcOp); - if (failed(typeExpander.expandFunctionSignature(funcOp, builder)) || - failed(typeExpander.expandAllReturnLikeTerminators<mlir::ReturnOp>( - funcOp, builder))) { - return signalPassFailure(); - } + if (failed(expandCallGraphTypes(getOperation(), typeExpander))) + signalPassFailure(); } }; @@ -79,14 +155,14 @@ // For any function which contains dynamic dims in its inputs or results, // rewrites it so that the dynamic dims are passed in/out. -std::unique_ptr<OperationPass<FuncOp>> createExpandFunctionDynamicDimsPass() { +std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionDynamicDimsPass() { return std::make_unique<Shape::ExpandFunctionDynamicDimsPass>(); } // For any function which contains ranked_shape argument/result types, // expands them to individual dynamic dimensions, inserting appropriate casts // within the function. -std::unique_ptr<OperationPass<FuncOp>> +std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionRankedShapeDimsPass() { return std::make_unique<Shape::ExpandFunctionRankedShapeDimsPass>(); }
diff --git a/iree/compiler/Dialect/Shape/Transforms/Passes.h b/iree/compiler/Dialect/Shape/Transforms/Passes.h index ab1544a..3cc2889 100644 --- a/iree/compiler/Dialect/Shape/Transforms/Passes.h +++ b/iree/compiler/Dialect/Shape/Transforms/Passes.h
@@ -18,12 +18,12 @@ // For any function which contains dynamic dims in its inputs or results, // rewrites it so that the dynamic dims are passed in/out. -std::unique_ptr<OperationPass<FuncOp>> createExpandFunctionDynamicDimsPass(); +std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionDynamicDimsPass(); // For any function which contains ranked_shape argument/result types, // expands them to individual dynamic dimensions, inserting appropriate casts // within the function. -std::unique_ptr<OperationPass<FuncOp>> +std::unique_ptr<OperationPass<ModuleOp>> createExpandFunctionRankedShapeDimsPass(); // Folds tensor.dim/memref.dim ops taking shape carrying ops as operands.
diff --git a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir index d0dfe66..be98e44 100644 --- a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir +++ b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_dynamic_dims.mlir
@@ -36,3 +36,18 @@ // CHECK: return %[[RESULT]], %[[SHAPE]] : tensor<?x3xf32>, !shapex.ranked_shape<[?,3]> return %1 : tensor<?x3xf32> } + +// ----- +// CHECK-LABEL: func @calls( +// CHECK-SAME: %[[ARG:.*]]: tensor<1x?x2x?xf32>, +// CHECK-SAME: %[[ARG_SHAPE:.*]]: !shapex.ranked_shape<[1,?,2,?]>) -> (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>) { +func @calls(%arg0 : tensor<1x?x2x?xf32>) -> tensor<1x?x2x?xf32> { + // CHECK: %[[ARG_TIED:.*]] = shapex.tie_shape %[[ARG]], %[[ARG_SHAPE]] : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]> + // CHECK: %[[ARG_TIED_SHAPE:.*]] = shapex.get_ranked_shape %[[ARG_TIED]] : tensor<1x?x2x?xf32> -> !shapex.ranked_shape<[1,?,2,?]> + // CHECK: %[[CALL_RESULT:.*]]:2 = call @calls(%[[ARG_TIED]], %[[ARG_TIED_SHAPE]]) : (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>) -> (tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]>) + %0 = std.call @calls(%arg0) : (tensor<1x?x2x?xf32>) -> tensor<1x?x2x?xf32> + // CHECK: %[[CALL_RESULT_TIED:.*]] = shapex.tie_shape %[[CALL_RESULT]]#0, %[[CALL_RESULT]]#1 : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]> + // CHECK: %[[CALL_RESULT_SHAPE:.*]] = shapex.get_ranked_shape %[[CALL_RESULT_TIED]] : tensor<1x?x2x?xf32> -> !shapex.ranked_shape<[1,?,2,?]> + // CHECK: return %[[CALL_RESULT_TIED]], %[[CALL_RESULT_SHAPE]] : tensor<1x?x2x?xf32>, !shapex.ranked_shape<[1,?,2,?]> + return %0 : tensor<1x?x2x?xf32> +}
diff --git a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir index 60e0cc4..9ad16d5 100644 --- a/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir +++ b/iree/compiler/Dialect/Shape/Transforms/test/expand_function_ranked_shape_dims.mlir
@@ -34,3 +34,12 @@ // CHECK: return %[[CAST_IDX0]], %[[CAST_IDX1]] return %rs : !shapex.ranked_shape<[?,?]> } + +// ----- +// CHECK-LABEL: @calls +// CHECK-SAME: (%arg0: index, %arg1: index) -> (index, index) +func @calls(%arg0 :!shapex.ranked_shape<[?,?]>) -> !shapex.ranked_shape<[?,?]> { + // CHECK: call @calls(%{{.*}}, %{{.*}}) : (index, index) -> (index, index) + %0 = std.call @calls(%arg0) : (!shapex.ranked_shape<[?,?]>) -> !shapex.ranked_shape<[?,?]> + return %0 : !shapex.ranked_shape<[?,?]> +}
diff --git a/iree/hal/cuda/executable_layout.c b/iree/hal/cuda/executable_layout.c index a89c7d9..21858ba 100644 --- a/iree/hal/cuda/executable_layout.c +++ b/iree/hal/cuda/executable_layout.c
@@ -15,6 +15,7 @@ typedef struct iree_hal_cuda_executable_layout_t { iree_hal_resource_t resource; iree_hal_cuda_context_wrapper_t* context; + iree_host_size_t push_constant_base_index; iree_host_size_t set_layout_count; iree_hal_descriptor_set_layout_t* set_layouts[]; } iree_hal_cuda_executable_layout_t; @@ -28,6 +29,21 @@ return (iree_hal_cuda_executable_layout_t*)base_value; } +static void iree_hal_cuda_executable_layout_destroy( + iree_hal_executable_layout_t* base_executable_layout) { + iree_hal_cuda_executable_layout_t* executable_layout = + iree_hal_cuda_executable_layout_cast(base_executable_layout); + iree_allocator_t host_allocator = executable_layout->context->host_allocator; + IREE_TRACE_ZONE_BEGIN(z0); + + for (iree_host_size_t i = 0; i < executable_layout->set_layout_count; ++i) { + iree_hal_descriptor_set_layout_release(executable_layout->set_layouts[i]); + } + iree_allocator_free(host_allocator, executable_layout); + + IREE_TRACE_ZONE_END(z0); +} + iree_status_t iree_hal_cuda_executable_layout_create( iree_hal_cuda_context_wrapper_t* context, iree_host_size_t set_layout_count, iree_hal_descriptor_set_layout_t** set_layouts, @@ -52,31 +68,20 @@ &executable_layout->resource); executable_layout->context = context; executable_layout->set_layout_count = set_layout_count; + iree_host_size_t binding_number = 0; for (iree_host_size_t i = 0; i < set_layout_count; ++i) { executable_layout->set_layouts[i] = set_layouts[i]; iree_hal_descriptor_set_layout_retain(set_layouts[i]); + binding_number += + iree_hal_cuda_descriptor_set_layout_binding_count(set_layouts[i]); } + executable_layout->push_constant_base_index = binding_number; *out_executable_layout = (iree_hal_executable_layout_t*)executable_layout; } IREE_TRACE_ZONE_END(z0); return status; } -static void iree_hal_cuda_executable_layout_destroy( - iree_hal_executable_layout_t* base_executable_layout) { - iree_hal_cuda_executable_layout_t* executable_layout = - iree_hal_cuda_executable_layout_cast(base_executable_layout); - iree_allocator_t host_allocator = executable_layout->context->host_allocator; - IREE_TRACE_ZONE_BEGIN(z0); - - for (iree_host_size_t i = 0; i < executable_layout->set_layout_count; ++i) { - iree_hal_descriptor_set_layout_release(executable_layout->set_layouts[i]); - } - iree_allocator_free(host_allocator, executable_layout); - - IREE_TRACE_ZONE_END(z0); -} - iree_host_size_t iree_hal_cuda_base_binding_index( iree_hal_executable_layout_t* base_executable_layout, uint32_t set) { iree_hal_cuda_executable_layout_t* executable_layout = @@ -91,6 +96,13 @@ return base_binding; } +iree_host_size_t iree_hal_cuda_push_constant_index( + iree_hal_executable_layout_t* base_executable_layout) { + iree_hal_cuda_executable_layout_t* executable_layout = + iree_hal_cuda_executable_layout_cast(base_executable_layout); + return executable_layout->push_constant_base_index; +} + const iree_hal_executable_layout_vtable_t iree_hal_cuda_executable_layout_vtable = { .destroy = iree_hal_cuda_executable_layout_destroy,
diff --git a/iree/hal/cuda/executable_layout.h b/iree/hal/cuda/executable_layout.h index 9b5a26e..adf29ea 100644 --- a/iree/hal/cuda/executable_layout.h +++ b/iree/hal/cuda/executable_layout.h
@@ -26,6 +26,10 @@ iree_host_size_t iree_hal_cuda_base_binding_index( iree_hal_executable_layout_t* executable_layout, uint32_t set); +// Return the base index for push constant data. +iree_host_size_t iree_hal_cuda_push_constant_index( + iree_hal_executable_layout_t* base_executable_layout); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus
diff --git a/iree/hal/cuda/graph_command_buffer.c b/iree/hal/cuda/graph_command_buffer.c index f5667fa..19ddb12 100644 --- a/iree/hal/cuda/graph_command_buffer.c +++ b/iree/hal/cuda/graph_command_buffer.c
@@ -37,6 +37,8 @@ } iree_hal_cuda_graph_command_buffer_t; #define IREE_HAL_CUDA_MAX_BINDING_COUNT 64 +// Kernel arguments contains binding and push constants. +#define IREE_HAL_CUDA_MAX_KERNEL_ARG 128 extern const iree_hal_command_buffer_vtable_t iree_hal_cuda_graph_command_buffer_vtable; @@ -62,8 +64,8 @@ "cuGraphCreate"); iree_hal_cuda_graph_command_buffer_t* command_buffer = NULL; size_t total_size = sizeof(*command_buffer) + - IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(void*) + - IREE_HAL_CUDA_MAX_BINDING_COUNT * sizeof(CUdeviceptr); + IREE_HAL_CUDA_MAX_KERNEL_ARG * sizeof(void*) + + IREE_HAL_CUDA_MAX_KERNEL_ARG * sizeof(CUdeviceptr); iree_status_t status = iree_allocator_malloc( context->host_allocator, total_size, (void**)&command_buffer); if (iree_status_is_ok(status)) { @@ -79,8 +81,8 @@ CUdeviceptr* device_ptrs = (CUdeviceptr*)(command_buffer->current_descriptor + - IREE_HAL_CUDA_MAX_BINDING_COUNT); - for (size_t i = 0; i < IREE_HAL_CUDA_MAX_BINDING_COUNT; i++) { + IREE_HAL_CUDA_MAX_KERNEL_ARG); + for (size_t i = 0; i < IREE_HAL_CUDA_MAX_KERNEL_ARG; i++) { command_buffer->current_descriptor[i] = &device_ptrs[i]; } @@ -330,8 +332,16 @@ iree_hal_command_buffer_t* base_command_buffer, iree_hal_executable_layout_t* executable_layout, iree_host_size_t offset, const void* values, iree_host_size_t values_length) { - return iree_make_status(IREE_STATUS_UNIMPLEMENTED, - "need cuda implementation"); + iree_hal_cuda_graph_command_buffer_t* command_buffer = + iree_hal_cuda_graph_command_buffer_cast(base_command_buffer); + iree_host_size_t constant_base_index = + iree_hal_cuda_push_constant_index(executable_layout) + + offset / sizeof(int32_t); + for (iree_host_size_t i = 0; i < values_length / sizeof(int32_t); i++) { + *((uint32_t*)command_buffer->current_descriptor[i + constant_base_index]) = + ((uint32_t*)values)[i]; + } + return iree_ok_status(); } // Tie together the binding index and its index in |bindings| array.
diff --git a/iree/test/e2e/regression/BUILD b/iree/test/e2e/regression/BUILD index 9aeca97..a034ab7 100644 --- a/iree/test/e2e/regression/BUILD +++ b/iree/test/e2e/regression/BUILD
@@ -19,6 +19,14 @@ ) BACKEND_TESTS = [ + "dynamic_abs.mlir", + "dynamic_add.mlir", + "dynamic_dot.mlir", + "dynamic_reduce_min.mlir", + "dynamic_torch_index_select_high_rank.mlir", + "dynamic_torch_index_select_negative.mlir", + "dynamic_torch_index_select_scalar.mlir", + "dynamic_torch_index_select_vector.mlir", "linalg_ops.mlir", ] @@ -26,14 +34,6 @@ name = "lit", srcs = enforce_glob( [ - "dynamic_abs.mlir", - "dynamic_add.mlir", - "dynamic_dot.mlir", - "dynamic_reduce_min.mlir", - "dynamic_torch_index_select_high_rank.mlir", - "dynamic_torch_index_select_negative.mlir", - "dynamic_torch_index_select_scalar.mlir", - "dynamic_torch_index_select_vector.mlir", "globals.mlir", "scalar.mlir", "tensor_cast.mlir", @@ -64,7 +64,7 @@ ) iree_check_single_backend_test_suite( - name = "check_regression_linalg_ops_dylib-llvm-aot", + name = "check_regression_dylib-llvm-aot", srcs = [ "linalg_ext_ops.mlir", "lowering_config.mlir", @@ -75,7 +75,7 @@ ) iree_check_single_backend_test_suite( - name = "check_regression_linalg_ops_vmvx", + name = "check_regression_vmvx", srcs = BACKEND_TESTS, compiler_flags = ["-iree-input-type=mhlo"], driver = "vmvx", @@ -83,9 +83,25 @@ ) iree_check_single_backend_test_suite( - name = "check_regression_linalg_ops_vulkan-spirv", + name = "check_regression_vulkan-spirv", srcs = BACKEND_TESTS, compiler_flags = ["-iree-input-type=mhlo"], driver = "vulkan", target_backend = "vulkan-spirv", ) + +iree_check_single_backend_test_suite( + name = "check_regression_cuda", + srcs = BACKEND_TESTS, + compiler_flags = ["-iree-input-type=mhlo"], + driver = "cuda", + tags = [ + # CUDA cuInit fails with sanitizer on. + "noasan", + "nomsan", + "notsan", + "noubsan", + "requires-gpu-nvidia", + ], + target_backend = "cuda", +)
diff --git a/iree/test/e2e/regression/CMakeLists.txt b/iree/test/e2e/regression/CMakeLists.txt index 38d6648..f812bb3 100644 --- a/iree/test/e2e/regression/CMakeLists.txt +++ b/iree/test/e2e/regression/CMakeLists.txt
@@ -14,14 +14,6 @@ NAME lit SRCS - "dynamic_abs.mlir" - "dynamic_add.mlir" - "dynamic_dot.mlir" - "dynamic_reduce_min.mlir" - "dynamic_torch_index_select_high_rank.mlir" - "dynamic_torch_index_select_negative.mlir" - "dynamic_torch_index_select_scalar.mlir" - "dynamic_torch_index_select_vector.mlir" "globals.mlir" "scalar.mlir" "tensor_cast.mlir" @@ -38,8 +30,16 @@ iree_check_single_backend_test_suite( NAME - check_regression_linalg_ops_dylib-llvm-aot + check_regression_dylib-llvm-aot SRCS + "dynamic_abs.mlir" + "dynamic_add.mlir" + "dynamic_dot.mlir" + "dynamic_reduce_min.mlir" + "dynamic_torch_index_select_high_rank.mlir" + "dynamic_torch_index_select_negative.mlir" + "dynamic_torch_index_select_scalar.mlir" + "dynamic_torch_index_select_vector.mlir" "linalg_ext_ops.mlir" "linalg_ops.mlir" "lowering_config.mlir" @@ -53,8 +53,16 @@ iree_check_single_backend_test_suite( NAME - check_regression_linalg_ops_vmvx + check_regression_vmvx SRCS + "dynamic_abs.mlir" + "dynamic_add.mlir" + "dynamic_dot.mlir" + "dynamic_reduce_min.mlir" + "dynamic_torch_index_select_high_rank.mlir" + "dynamic_torch_index_select_negative.mlir" + "dynamic_torch_index_select_scalar.mlir" + "dynamic_torch_index_select_vector.mlir" "linalg_ops.mlir" TARGET_BACKEND "vmvx" @@ -66,8 +74,16 @@ iree_check_single_backend_test_suite( NAME - check_regression_linalg_ops_vulkan-spirv + check_regression_vulkan-spirv SRCS + "dynamic_abs.mlir" + "dynamic_add.mlir" + "dynamic_dot.mlir" + "dynamic_reduce_min.mlir" + "dynamic_torch_index_select_high_rank.mlir" + "dynamic_torch_index_select_negative.mlir" + "dynamic_torch_index_select_scalar.mlir" + "dynamic_torch_index_select_vector.mlir" "linalg_ops.mlir" TARGET_BACKEND "vulkan-spirv" @@ -77,4 +93,31 @@ "-iree-input-type=mhlo" ) +iree_check_single_backend_test_suite( + NAME + check_regression_cuda + SRCS + "dynamic_abs.mlir" + "dynamic_add.mlir" + "dynamic_dot.mlir" + "dynamic_reduce_min.mlir" + "dynamic_torch_index_select_high_rank.mlir" + "dynamic_torch_index_select_negative.mlir" + "dynamic_torch_index_select_scalar.mlir" + "dynamic_torch_index_select_vector.mlir" + "linalg_ops.mlir" + TARGET_BACKEND + "cuda" + DRIVER + "cuda" + COMPILER_FLAGS + "-iree-input-type=mhlo" + LABELS + "noasan" + "nomsan" + "notsan" + "noubsan" + "requires-gpu-nvidia" +) + ### BAZEL_TO_CMAKE_PRESERVES_ALL_CONTENT_BELOW_THIS_LINE ###
diff --git a/iree/test/e2e/regression/dynamic_abs.mlir b/iree/test/e2e/regression/dynamic_abs.mlir index e1c6343..fa63d0d 100644 --- a/iree/test/e2e/regression/dynamic_abs.mlir +++ b/iree/test/e2e/regression/dynamic_abs.mlir
@@ -1,12 +1,8 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot %s | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx %s | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv %s | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @dynamic_tensor -func @dynamic_tensor() -> tensor<?x?xf32> { +func @dynamic_tensor() { %input = util.dynamic_shape_constant dense<[[-1.0, 2.0, -3.0], [4.0, -5.0, 6.0]]> : tensor<2x3xf32> -> tensor<?x?xf32> %res = "mhlo.abs"(%input) : (tensor<?x?xf32>) -> tensor<?x?xf32> - return %res : tensor<?x?xf32> + %dshape = util.do_not_optimize(%res) : tensor<?x?xf32> + %result = tensor.cast %dshape : tensor<?x?xf32> to tensor<2x3xf32> + check.expect_almost_eq_const(%result, dense<[[1.0, 2.0, 3.0],[4.0, 5.0, 6.0]]> : tensor<2x3xf32>) : tensor<2x3xf32> + return } - -// CHECK: 2x3xf32=[1 2 3][4 5 6]
diff --git a/iree/test/e2e/regression/dynamic_add.mlir b/iree/test/e2e/regression/dynamic_add.mlir index b9bb49d..ab5ee2c 100644 --- a/iree/test/e2e/regression/dynamic_add.mlir +++ b/iree/test/e2e/regression/dynamic_add.mlir
@@ -1,11 +1,9 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" -function-input="2x4xf32=[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]" %s | IreeFileCheck %s) - -// CHECK: EXEC @main -// CHECK: 2x4xf32=[6 8 10 12][-6 -8 -10 -12] - -func @main(%arg0: tensor<?x4xf32>, %arg1: tensor<?x4xf32>) -> tensor<?x4xf32> { - %0 = "mhlo.add"(%arg0, %arg1) : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32> - return %0: tensor<?x4xf32> +func @main() { + %lhs = util.dynamic_shape_constant dense<[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]> : tensor<2x4xf32> -> tensor<?x4xf32> + %rhs = util.dynamic_shape_constant dense<[[5.0,6.0,7.0,8.0],[-5.0,-6.0,-7.0,-8.0]]> : tensor<2x4xf32> -> tensor<?x4xf32> + %2 = "mhlo.add"(%lhs, %rhs) : (tensor<?x4xf32>, tensor<?x4xf32>) -> tensor<?x4xf32> + %3 = util.do_not_optimize(%2) : tensor<?x4xf32> + %result = tensor.cast %3 : tensor<?x4xf32> to tensor<2x4xf32> + check.expect_almost_eq_const(%result, dense<[[6.0, 8.0, 10.0, 12.0],[-6.0, -8.0, -10.0, -12.0]]> : tensor<2x4xf32>) : tensor<2x4xf32> + return }
diff --git a/iree/test/e2e/regression/dynamic_dot.mlir b/iree/test/e2e/regression/dynamic_dot.mlir index e9d82d7..c60d2b2 100644 --- a/iree/test/e2e/regression/dynamic_dot.mlir +++ b/iree/test/e2e/regression/dynamic_dot.mlir
@@ -1,9 +1,4 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot %s | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx %s | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv %s | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @dynamic_dot -func @dynamic_dot() -> tensor<?x?xf32> { +func @dynamic_dot() { %lhs = util.dynamic_shape_constant dense<[ [15.0, 14.0, 13.0], [12.0, 11.0, 10.0], @@ -15,7 +10,13 @@ [10.0, 09.0, 08.0, 07.0, 06.0], [05.0, 04.0, 03.0, 02.0, 01.0]]> : tensor<3x5xf32> -> tensor<?x?xf32> %res = "mhlo.dot"(%lhs, %rhs) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32> - return %res : tensor<?x?xf32> + %dshape = util.do_not_optimize(%res) : tensor<?x?xf32> + %result = tensor.cast %dshape : tensor<?x?xf32> to tensor<5x5xf32> + check.expect_almost_eq_const(%result, + dense<[[430.0, 388.0, 346.0, 304.0, 262.0], + [340.0, 307.0, 274.0, 241.0, 208.0], + [250.0, 226.0, 202.0, 178.0, 154.0], + [160.0, 145.0, 130.0, 115.0, 100.0], + [70.0, 64.0, 58.0, 52.0, 46.0]]> : tensor<5x5xf32>) : tensor<5x5xf32> + return } - -// CHECK: 5x5xf32=[430 388 346 304 262][340 307 274 241 208][250 226 202 178 154][160 145 130 115 100][70 64 58 52 46]
diff --git a/iree/test/e2e/regression/dynamic_reduce_min.mlir b/iree/test/e2e/regression/dynamic_reduce_min.mlir index cb5b0ae..b912c84 100644 --- a/iree/test/e2e/regression/dynamic_reduce_min.mlir +++ b/iree/test/e2e/regression/dynamic_reduce_min.mlir
@@ -1,16 +1,12 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x4xf32=[[1.0,2.0,3.0,4.0],[-1.0,-2.0,-3.0,-4.0]]" %s | IreeFileCheck %s) - -func @reduce_min(%arg0: tensor<?x?xf32>) -> tensor<f32> { +func @reduce_min() { + %input = util.dynamic_shape_constant + dense<[[1.0, 2.0, 3.0, 4.0],[-1.0 ,-2.0 ,-3.0 ,-4.0]]> : tensor<2x4xf32> -> tensor<?x?xf32> %0 = mhlo.constant dense<0x7F800000> : tensor<f32> - %1 = "mhlo.reduce"(%arg0, %0) ( { + %1 = "mhlo.reduce"(%input, %0) ( { ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>): %2 = mhlo.minimum %arg1, %arg2 : tensor<f32> "mhlo.return"(%2) : (tensor<f32>) -> () }) {dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<?x?xf32>, tensor<f32>) -> tensor<f32> - return %1 : tensor<f32> + check.expect_almost_eq_const(%1, dense<-4.0> : tensor<f32>) : tensor<f32> + return } - -// CHECK: EXEC @reduce_min -// CHECK: f32=-4
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir index 570f8a5..b15c6a3 100644 --- a/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir +++ b/iree/test/e2e/regression/dynamic_torch_index_select_high_rank.mlir
@@ -1,59 +1,31 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x2xi32=[6,7][8,9]" -function-input="2x2x2x2xi32=[[[0,1][1,0]][[0,0][1,1]]][[[1,1][0,0]][[0,1][1,0]]]" | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @torch_index_select1 -func @torch_index_select1(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 1 : i64, dim = 1 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> - return %0 : tensor<?x?x?x?xi32> +func @torch_index_select1() { + %lhs = util.dynamic_shape_constant dense<[[6,7],[8,9]]> : tensor<2x2xi32> -> tensor<?x?xi32> + %rhs = util.dynamic_shape_constant dense<[[[[0,1],[1,0]],[[0,0],[1,1]]],[[[1,1],[0,0]],[[0,1],[1,0]]]]> : tensor<2x2x2x2xi32> -> tensor<?x?x?x?xi32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 1 : i64, dim = 1 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?x?x?xi32> + %result = tensor.cast %dshape : tensor<?x?x?x?xi32> to tensor<2x2x2x2xi32> + check.expect_eq_const(%result, + dense<[[[[6, 7],[7, 6]], + [[6, 6],[7, 7]]], + [[[9, 9],[8, 8]], + [[8, 9],[9, 8]]]]> : tensor<2x2x2x2xi32>) : tensor<2x2x2x2xi32> + return } -// CHECK: 2x2x2x2xi32=[ -// CHECK-SAME: [ -// CHECK-SAME: [6 7][7 6] -// CHECK-SAME: ][ -// CHECK-SAME: [6 6][7 7] -// CHECK-SAME: ] -// CHECK-SAME: ][ -// CHECK-SAME: [ -// CHECK-SAME: [9 9][8 8] -// CHECK-SAME: ][ -// CHECK-SAME: [8 9][9 8] -// CHECK-SAME: ] -// CHECK-SAME: ] - -// CHECK-LABEL: EXEC @torch_index_select2 -func @torch_index_select2(%arg0: tensor<?x?xi32>, %arg1: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32> - return %0 : tensor<?x?x?x?x?xi32> +func @torch_index_select2() { + %lhs = util.dynamic_shape_constant dense<[[6,7],[8,9]]> : tensor<2x2xi32> -> tensor<?x?xi32> + %rhs = util.dynamic_shape_constant dense<[[[[0,1],[1,0]],[[0,0],[1,1]]],[[[1,1],[0,0]],[[0,1],[1,0]]]]> : tensor<2x2x2x2xi32> -> tensor<?x?x?x?xi32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?xi32>, tensor<?x?x?x?xi32>) -> tensor<?x?x?x?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?x?x?x?xi32> + %result = tensor.cast %dshape : tensor<?x?x?x?x?xi32> to tensor<2x2x2x2x2xi32> + check.expect_eq_const(%result, + dense<[[[[[6, 7],[8, 9]], + [[8, 9],[6, 7]]], + [[[6, 7],[6, 7]], + [[8, 9],[8, 9]]]], + [[[[8, 9],[8, 9]], + [[6, 7],[6, 7]]], + [[[6, 7],[8, 9]], + [[8, 9],[6, 7]]]]]> : tensor<2x2x2x2x2xi32>) : tensor<2x2x2x2x2xi32> + return } - -// CHECK: 2x2x2x2x2xi32=[ -// CHECK-SAME: [ -// CHECK-SAME: [ -// CHECK-SAME: [6 7][8 9] -// CHECK-SAME: ][ -// CHECK-SAME: [8 9][6 7] -// CHECK-SAME: ] -// CHECK-SAME: ][ -// CHECK-SAME: [ -// CHECK-SAME: [6 7][6 7] -// CHECK-SAME: ][ -// CHECK-SAME: [8 9][8 9] -// CHECK-SAME: ] -// CHECK-SAME: ] -// CHECK-SAME: ][ -// CHECK-SAME: [ -// CHECK-SAME: [ -// CHECK-SAME: [8 9][8 9] -// CHECK-SAME: ][ -// CHECK-SAME: [6 7][6 7] -// CHECK-SAME: ] -// CHECK-SAME: ][ -// CHECK-SAME: [ -// CHECK-SAME: [6 7][8 9] -// CHECK-SAME: ][ -// CHECK-SAME: [8 9][6 7] -// CHECK-SAME: ] -// CHECK-SAME: ] -// CHECK-SAME: ]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir index f23dae5..19f8367 100644 --- a/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir +++ b/iree/test/e2e/regression/dynamic_torch_index_select_negative.mlir
@@ -1,15 +1,11 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="2x2x2xi32=[[100, 101] [110, 111]] [[200, 201] [210, 211]]" -function-input="2x2x2xi32=[[0, 1] [1, 0]] [[0, 0] [1, 1]]" | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @torch_index_select1 -func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?x?x?xi32>) -> tensor<?x?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = -1 : i64, dim = -1 : i64} : (tensor<?x?x?xi32>, tensor<?x?x?xi32>) -> tensor<?x?x?xi32> - return %0 : tensor<?x?x?xi32> +func @torch_index_select1() { + %lhs = util.dynamic_shape_constant dense<[[[100, 101],[110, 111]],[[200, 201],[210, 211]]]> : tensor<2x2x2xi32> -> tensor<?x?x?xi32> + %rhs = util.dynamic_shape_constant dense<[[[0, 1],[1, 0]],[[0, 0],[1, 1]]]> : tensor<2x2x2xi32> -> tensor<?x?x?xi32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = -1 : i64, dim = -1 : i64} : (tensor<?x?x?xi32>, tensor<?x?x?xi32>) -> tensor<?x?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32> + %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<2x2x2xi32> + check.expect_eq_const(%result, + dense<[[[100, 101],[111, 110]], + [[200, 200],[211, 211]]]> : tensor<2x2x2xi32>) : tensor<2x2x2xi32> + return } - -// CHECK: 2x2x2xi32=[ -// CHECK-SAME: [100 101][111 110] -// CHECK-SAME: ][ -// CHECK-SAME: [200 200][211 211] -// CHECK-SAME: ]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir index e97cf8f..f68962d 100644 --- a/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir +++ b/iree/test/e2e/regression/dynamic_torch_index_select_scalar.mlir
@@ -1,19 +1,35 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="5x1x5xi32=[[1,2,3,4,5]] [[6,7,8,9,10]] [[11,12,13,14,15]] [[16,17,18,19,20]] [[21,22,23,24,25]]" -function-input="i32=0" | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @torch_index_select1 -func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<i32>) -> tensor<?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32> - return %0 : tensor<?x?xi32> +func @torch_index_select1() { + %lhs = util.dynamic_shape_constant + dense<[[[1,2,3,4,5]], + [[6,7,8,9,10]], + [[11,12,13,14,15]], + [[16,17,18,19,20]], + [[21,22,23,24,25]]]> : tensor<5x1x5xi32> -> tensor<?x?x?xi32> + %rhs = util.unfoldable_constant dense<0> : tensor<i32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?xi32> + %result = tensor.cast %dshape : tensor<?x?xi32> to tensor<1x5xi32> + check.expect_eq_const(%result, + dense<[[1, 2, 3, 4, 5]]> : tensor<1x5xi32>) : tensor<1x5xi32> + return } -// CHECK: 1x5xi32=[1 2 3 4 5] - -// CHECK-LABEL: EXEC @torch_index_select2 -func @torch_index_select2(%arg0: tensor<?x?x?xi32>, %arg1: tensor<i32>) -> tensor<?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32> - return %0 : tensor<?x?xi32> +func @torch_index_select2() { + %lhs = util.dynamic_shape_constant + dense<[[[1,2,3,4,5]], + [[6,7,8,9,10]], + [[11,12,13,14,15]], + [[16,17,18,19,20]], + [[21,22,23,24,25]]]> : tensor<5x1x5xi32> -> tensor<?x?x?xi32> + %rhs = util.unfoldable_constant dense<0> : tensor<i32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<i32>) -> tensor<?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?xi32> + %result = tensor.cast %dshape : tensor<?x?xi32> to tensor<5x5xi32> + check.expect_eq_const(%result, + dense<[[1, 2, 3, 4, 5], + [6, 7, 8, 9, 10], + [11, 12, 13, 14, 15], + [16, 17, 18, 19, 20], + [21, 22, 23, 24, 25]]> : tensor<5x5xi32>) : tensor<5x5xi32> + return } - -// CHECK: 5x5xi32=[1 2 3 4 5][6 7 8 9 10][11 12 13 14 15][16 17 18 19 20][21 22 23 24 25]
diff --git a/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir b/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir index 64e13dd..1ef9007 100644 --- a/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir +++ b/iree/test/e2e/regression/dynamic_torch_index_select_vector.mlir
@@ -1,29 +1,26 @@ -// RUN: [[ $IREE_LLVMAOT_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=dylib-llvm-aot -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s) -// RUN: [[ $IREE_VMVX_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vmvx -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s) -// RUN: [[ $IREE_VULKAN_DISABLE == 1 ]] || (iree-run-mlir %s --iree-input-type=mhlo -iree-hal-target-backends=vulkan-spirv -function-input="3x2x2xi32=[[1, 2] [3, 4]] [[5, 6] [7, 8]] [[9, 10] [11, 12]]" -function-input="2xi32=[0, 1]" | IreeFileCheck %s) - -// CHECK-LABEL: EXEC @torch_index_select1 -func @torch_index_select1(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?xi32>) -> tensor<?x?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32> - return %0 : tensor<?x?x?xi32> +func @torch_index_select1() { + %lhs = util.dynamic_shape_constant + dense<[[[1, 2],[3, 4]],[[5, 6],[7, 8]],[[9, 10],[11, 12]]]> : tensor<3x2x2xi32> -> tensor<?x?x?xi32> + %rhs = util.dynamic_shape_constant dense<[0, 1]> : tensor<2xi32> -> tensor<?xi32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 1 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32> + %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<3x2x2xi32> + check.expect_eq_const(%result, + dense<[[[1, 2],[3, 4]], + [[5, 6],[7, 8]], + [[9, 10],[11, 12]]]> : tensor<3x2x2xi32>) : tensor<3x2x2xi32> + return } -// CHECK: 3x2x2xi32=[ -// CHECK-SAME: [1 2][3 4] -// CHECK-SAME: ][ -// CHECK-SAME: [5 6][7 8] -// CHECK-SAME: ][ -// CHECK-SAME: [9 10][11 12] -// CHECK-SAME: ] - -// CHECK-LABEL: EXEC @torch_index_select2 -func @torch_index_select2(%arg0: tensor<?x?x?xi32>, %arg1: tensor<?xi32>) -> tensor<?x?x?xi32> { - %0 = "mhlo.torch_index_select"(%arg0, %arg1) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32> - return %0 : tensor<?x?x?xi32> +func @torch_index_select2() { + %lhs = util.dynamic_shape_constant + dense<[[[1, 2],[3, 4]],[[5, 6],[7, 8]],[[9, 10],[11, 12]]]> : tensor<3x2x2xi32> -> tensor<?x?x?xi32> + %rhs = util.dynamic_shape_constant dense<[0, 1]> : tensor<2xi32> -> tensor<?xi32> + %0 = "mhlo.torch_index_select"(%lhs, %rhs) {batch_dims = 0 : i64, dim = 0 : i64} : (tensor<?x?x?xi32>, tensor<?xi32>) -> tensor<?x?x?xi32> + %dshape = util.do_not_optimize(%0) : tensor<?x?x?xi32> + %result = tensor.cast %dshape : tensor<?x?x?xi32> to tensor<2x2x2xi32> + check.expect_eq_const(%result, + dense<[[[1, 2],[3, 4]], + [[5, 6],[7, 8]]]> : tensor<2x2x2xi32>) : tensor<2x2x2xi32> + return } - -// CHECK: 2x2x2xi32=[ -// CHECK-SAME: [1 2][3 4] -// CHECK-SAME: ][ -// CHECK-SAME: [5 6][7 8] -// CHECK-SAME: ]
diff --git a/iree/test/e2e/xla_ops/BUILD b/iree/test/e2e/xla_ops/BUILD index b9175a1..9a4e62b 100644 --- a/iree/test/e2e/xla_ops/BUILD +++ b/iree/test/e2e/xla_ops/BUILD
@@ -62,8 +62,10 @@ "remainder.mlir", "reshape.mlir", "reverse.mlir", + "rng_uniform.mlir", "rsqrt.mlir", "scatter.mlir", + "scatter_dynamic.mlir", "select.mlir", "sine.mlir", "slice.mlir", @@ -78,7 +80,6 @@ include = ["*.mlir"], exclude = [ "round.mlir", - "scatter_dynamic.mlir", ], ), compiler_flags = ["-iree-input-type=mhlo"], @@ -207,6 +208,7 @@ "remainder.mlir", "reshape.mlir", "reverse.mlir", + "rng_uniform.mlir", "rsqrt.mlir", "scatter.mlir", "scatter_dynamic.mlir", @@ -291,6 +293,7 @@ "scatter.mlir", # TODO(GH-6601): Enable the test. "scatter_dynamic.mlir", # TODO(GH-6601): Enable the test. "sort.mlir", + "rng_uniform.mlir", # TODO(GH-6995): Enable the test. ], ), compiler_flags = ["-iree-input-type=mhlo"], @@ -340,6 +343,7 @@ "remainder.mlir", "reshape.mlir", "reverse.mlir", + "rng_uniform.mlir", "rsqrt.mlir", "scatter.mlir", "scatter_dynamic.mlir",
diff --git a/iree/test/e2e/xla_ops/CMakeLists.txt b/iree/test/e2e/xla_ops/CMakeLists.txt index 3a7007a..392dfa4 100644 --- a/iree/test/e2e/xla_ops/CMakeLists.txt +++ b/iree/test/e2e/xla_ops/CMakeLists.txt
@@ -52,8 +52,10 @@ "remainder.mlir" "reshape.mlir" "reverse.mlir" + "rng_uniform.mlir" "rsqrt.mlir" "scatter.mlir" + "scatter_dynamic.mlir" "select.mlir" "sine.mlir" "slice.mlir" @@ -120,6 +122,7 @@ "remainder.mlir" "reshape.mlir" "reverse.mlir" + "rng_uniform.mlir" "rsqrt.mlir" "scatter.mlir" "scatter_dynamic.mlir" @@ -241,6 +244,7 @@ "remainder.mlir" "reshape.mlir" "reverse.mlir" + "rng_uniform.mlir" "rsqrt.mlir" "scatter.mlir" "scatter_dynamic.mlir"
diff --git a/iree/test/e2e/xla_ops/rng_uniform.mlir b/iree/test/e2e/xla_ops/rng_uniform.mlir new file mode 100644 index 0000000..7e496db --- /dev/null +++ b/iree/test/e2e/xla_ops/rng_uniform.mlir
@@ -0,0 +1,34 @@ +// Note that they are stateless random generators, so they have fixed results. +func @rng_uniform_1d() { + %min = util.unfoldable_constant dense<-10.0> : tensor<f32> + %max = util.unfoldable_constant dense<10.0> : tensor<f32> + %shape = util.unfoldable_constant dense<[10]> : tensor<1xi32> + %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<1xi32>) -> tensor<10xf32> + check.expect_almost_eq_const(%res, dense<[ + -9.99994, -4.8613, 0.277344, 5.41599, -9.44537, -4.30673, 0.831918, 5.97056, -8.8908, -3.75215 + ]> : tensor<10xf32>) : tensor<10xf32> + return +} + +func @rng_uniform_2d() { + %min = util.unfoldable_constant dense<-10.0> : tensor<f32> + %max = util.unfoldable_constant dense<10.0> : tensor<f32> + %shape = util.unfoldable_constant dense<[3, 3]> : tensor<2xi32> + %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<2xi32>) -> tensor<3x3xf32> + check.expect_almost_eq_const(%res, dense<[ + [6.55154, -8.30982, -3.17117], + [1.75741, 6.89606, -7.9653], + [-3.03671, 2.10193, 7.24057]]> : tensor<3x3xf32>) : tensor<3x3xf32> + return +} + +func @rng_uniform_3d() { + %min = util.unfoldable_constant dense<-10.0> : tensor<f32> + %max = util.unfoldable_constant dense<10.0> : tensor<f32> + %shape = util.unfoldable_constant dense<[2, 2, 2]> : tensor<3xi32> + %res = "mhlo.rng_uniform"(%min, %max, %shape) : (tensor<f32>, tensor<f32>, tensor<3xi32>) -> tensor<2x2x2xf32> + check.expect_almost_eq_const(%res, dense<[ + [[3.04814, 8.18679], [-1.74598, 3.39266]], + [[-6.91349, -1.77484], [8.29239, -6.56897]]]> : tensor<2x2x2xf32>) : tensor<2x2x2xf32> + return +}