| //===- InferIntRangeInterfaceImpls.cpp - Integer range impls for gpu -===// |
| // |
| // Part of the LLVM Project, 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 |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "mlir/Dialect/GPU/IR/GPUDialect.h" |
| #include "mlir/IR/Matchers.h" |
| #include "mlir/Interfaces/FunctionInterfaces.h" |
| #include "mlir/Interfaces/InferIntRangeInterface.h" |
| #include <optional> |
| |
| using namespace mlir; |
| using namespace mlir::gpu; |
| |
| // Maximum grid and block dimensions of all known GPUs are less than 2^32. |
| static constexpr uint64_t kMaxDim = std::numeric_limits<uint32_t>::max(); |
| // Maximum cluster size. |
| static constexpr uint64_t kMaxClusterDim = 16; |
| // Maximum subgroups are no larger than 128. |
| static constexpr uint64_t kMaxSubgroupSize = 128; |
| |
| static ConstantIntRanges getIndexRange(uint64_t umin, uint64_t umax) { |
| unsigned width = IndexType::kInternalStorageBitWidth; |
| return ConstantIntRanges::fromUnsigned(APInt(width, umin), |
| APInt(width, umax)); |
| } |
| |
| static uint64_t zext(uint32_t arg) { return static_cast<uint64_t>(arg); } |
| |
| static Value valueByDim(KernelDim3 dims, Dimension dim) { |
| switch (dim) { |
| case Dimension::x: |
| return dims.x; |
| case Dimension::y: |
| return dims.y; |
| case Dimension::z: |
| return dims.z; |
| } |
| llvm_unreachable("All dimension enum cases handled above"); |
| } |
| |
| static std::optional<uint32_t> |
| getKnownLaunchAttr(GPUFuncOp func, DimensionKind dims, Dimension dim) { |
| DenseI32ArrayAttr bounds; |
| switch (dims) { |
| case DimensionKind::Other: |
| return std::nullopt; |
| case DimensionKind::Block: |
| bounds = func.getKnownBlockSizeAttr(); |
| break; |
| case DimensionKind::Grid: |
| bounds = func.getKnownGridSizeAttr(); |
| break; |
| case DimensionKind::Cluster: |
| bounds = func.getKnownClusterSizeAttr(); |
| break; |
| } |
| if (!bounds) |
| return std::nullopt; |
| if (bounds.size() <= static_cast<uint32_t>(dim)) |
| return std::nullopt; |
| return bounds[static_cast<uint32_t>(dim)]; |
| } |
| |
| static std::optional<uint32_t> getKnownLaunchAttr(FunctionOpInterface func, |
| StringRef attrName, |
| Dimension dim) { |
| auto bounds = func.getOperation()->getAttrOfType<DenseI32ArrayAttr>(attrName); |
| if (!bounds) |
| return std::nullopt; |
| if (bounds.size() <= static_cast<uint32_t>(dim)) |
| return std::nullopt; |
| return bounds[static_cast<uint32_t>(dim)]; |
| } |
| |
| std::optional<uint32_t> |
| mlir::gpu::getKnownDimensionSizeAround(Operation *op, DimensionKind kind, |
| Dimension dim) { |
| if (auto launch = op->getParentOfType<LaunchOp>()) { |
| KernelDim3 bounds; |
| switch (kind) { |
| case DimensionKind::Other: |
| return std::nullopt; |
| case DimensionKind::Block: |
| bounds = launch.getBlockSizeOperandValues(); |
| break; |
| case DimensionKind::Grid: |
| bounds = launch.getGridSizeOperandValues(); |
| break; |
| case DimensionKind::Cluster: |
| if (launch.hasClusterSize()) { |
| auto clusterBounds = launch.getClusterSizeOperandValues(); |
| if (clusterBounds) |
| bounds = *clusterBounds; |
| } |
| break; |
| } |
| Value maybeBound = valueByDim(bounds, dim); |
| APInt value; |
| if (maybeBound && matchPattern(maybeBound, m_ConstantInt(&value))) |
| return value.getZExtValue(); |
| } |
| |
| if (auto gpuFunc = op->getParentOfType<GPUFuncOp>()) { |
| auto inherentAttr = getKnownLaunchAttr(gpuFunc, kind, dim); |
| if (inherentAttr) |
| return inherentAttr; |
| } |
| if (auto func = op->getParentOfType<FunctionOpInterface>()) { |
| StringRef attrName; |
| switch (kind) { |
| case DimensionKind::Other: |
| return std::nullopt; |
| case DimensionKind::Block: |
| attrName = GPUDialect::KnownBlockSizeAttrHelper::getNameStr(); |
| break; |
| case DimensionKind::Grid: |
| attrName = GPUDialect::KnownGridSizeAttrHelper::getNameStr(); |
| break; |
| case DimensionKind::Cluster: |
| attrName = GPUDialect::KnownClusterSizeAttrHelper::getNameStr(); |
| break; |
| } |
| auto discardableAttr = getKnownLaunchAttr(func, attrName, dim); |
| if (discardableAttr) |
| return discardableAttr; |
| } |
| return std::nullopt; |
| } |
| |
| void ClusterDimOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void ClusterDimBlocksOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| if (auto known = getKnownDimensionSizeAround(*this, DimensionKind::Cluster, |
| getDimension())) |
| return setResultRange(getResult(), |
| getIndexRange(zext(*known), zext(*known))); |
| |
| uint64_t max = kMaxClusterDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void ClusterIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void ClusterBlockIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxClusterDim; |
| if (auto known = getKnownDimensionSizeAround(*this, DimensionKind::Cluster, |
| getDimension())) |
| max = zext(*known); |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void BlockDimOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| std::optional<uint32_t> knownVal = |
| getKnownDimensionSizeAround(*this, DimensionKind::Block, getDimension()); |
| if (knownVal) |
| return setResultRange(getResult(), |
| getIndexRange(zext(*knownVal), zext(*knownVal))); |
| |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void BlockIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto fromContext = getKnownDimensionSizeAround(*this, DimensionKind::Grid, |
| getDimension())) |
| max = zext(*fromContext); |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void GridDimOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| std::optional<uint32_t> knownVal = |
| getKnownDimensionSizeAround(*this, DimensionKind::Grid, getDimension()); |
| if (knownVal) |
| return setResultRange(getResult(), |
| getIndexRange(zext(*knownVal), zext(*knownVal))); |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void ThreadIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto fromContext = getKnownDimensionSizeAround( |
| *this, DimensionKind::Block, getDimension())) |
| max = zext(*fromContext); |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void LaneIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxSubgroupSize; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void SubgroupIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(0, max - 1ULL)); |
| } |
| |
| void GlobalIdOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| if (auto specified = getUpperBound()) |
| return setResultRange(getResult(), |
| getIndexRange(0, specified->getZExtValue() - 1ULL)); |
| |
| uint64_t blockDimMax = zext( |
| getKnownDimensionSizeAround(*this, DimensionKind::Block, getDimension()) |
| .value_or(kMaxDim)); |
| uint64_t gridDimMax = zext( |
| getKnownDimensionSizeAround(*this, DimensionKind::Grid, getDimension()) |
| .value_or(kMaxDim)); |
| setResultRange(getResult(), |
| getIndexRange(0, (blockDimMax * gridDimMax) - 1ULL)); |
| } |
| |
| void NumSubgroupsOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxDim; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void SubgroupSizeOp::inferResultRanges(ArrayRef<ConstantIntRanges>, |
| SetIntRangeFn setResultRange) { |
| uint64_t max = kMaxSubgroupSize; |
| if (auto specified = getUpperBound()) |
| max = specified->getZExtValue(); |
| setResultRange(getResult(), getIndexRange(1, max)); |
| } |
| |
| void LaunchOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges, |
| SetIntRangeFn setResultRange) { |
| auto setRange = [&](const ConstantIntRanges &argRange, Value dimResult, |
| Value idxResult) { |
| if (argRange.umin().getBitWidth() != IndexType::kInternalStorageBitWidth) |
| return; |
| ConstantIntRanges dimRange = |
| argRange.intersection(getIndexRange(1, kMaxDim)); |
| setResultRange(dimResult, dimRange); |
| ConstantIntRanges idxRange = |
| getIndexRange(0, dimRange.umax().getZExtValue() - 1); |
| setResultRange(idxResult, idxRange); |
| }; |
| |
| argRanges = argRanges.drop_front(getAsyncDependencies().size()); |
| KernelDim3 gridDims = getGridSize(); |
| KernelDim3 blockIds = getBlockIds(); |
| setRange(argRanges[0], gridDims.x, blockIds.x); |
| setRange(argRanges[1], gridDims.y, blockIds.y); |
| setRange(argRanges[2], gridDims.z, blockIds.z); |
| KernelDim3 blockDims = getBlockSize(); |
| KernelDim3 threadIds = getThreadIds(); |
| setRange(argRanges[3], blockDims.x, threadIds.x); |
| setRange(argRanges[4], blockDims.y, threadIds.y); |
| setRange(argRanges[5], blockDims.z, threadIds.z); |
| } |
