[mlir][GPU] Refactor, improve constant size information handling

mlir/include/mlir/Dialect/GPU/IR/GPUBase.td

@@ -101,6 +101,27 @@ def GPU_AddressSpaceAttr :
 
 def GPU_AddressSpaceAttrArray : TypedArrayAttrBase<GPU_AddressSpaceAttr, "GPU Address Space array">;
 
+def GPU_Dimension : GPU_I32Enum<"Dimension",
+    "a dimension, either 'x', 'y', or 'z'",
+    [
+      I32EnumAttrCase<"x", 0>,
+      I32EnumAttrCase<"y", 1>,
+      I32EnumAttrCase<"z", 2>
+    ]>;
+
+def GPU_DimensionAttr : EnumAttr<GPU_Dialect, GPU_Dimension, "dim">;
+
+def GPU_DimensionKind : I32Enum<"DimensionKind",
+  "the possible kinds of launch dimension",
+  [
+    I32EnumCase<"Other", 0, "other">,
+    I32EnumCase<"Block", 1, "block">,
+    I32EnumCase<"Grid", 2, "grid">,
+    I32EnumCase<"Cluster", 3, "cluster">
+  ]> {
+  let cppNamespace = "::mlir::gpu";
+}
+
 //===----------------------------------------------------------------------===//
 // GPU Types.
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h

@@ -221,4 +221,12 @@ class SparseSpGEMMOpHandleType
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/IR/GPUOps.h.inc"
 
+namespace mlir::gpu {
+/// Retrieve the constant bounds for a given dimension and dimension kind
+/// from the context surrounding `op`, if known, and return them. This will
+/// check the bounds on an enclosing `gpu.launch`, an enclosing `gpu.func`, and
+/// any `gpu.known_*_size` on other function-like operations, in that order.
+std::optional<uint32_t>
+getKnownDimensionSizeAround(Operation *op, DimensionKind kind, Dimension dim);
+} // namespace mlir::gpu
 #endif // MLIR_DIALECT_GPU_IR_GPUDIALECT_H

mlir/include/mlir/Dialect/GPU/IR/GPUOps.td

@@ -38,18 +38,6 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 class GPU_Op<string mnemonic, list<Trait> traits = []> :
     Op<GPU_Dialect, mnemonic, traits>;
 
-def GPU_Dimension : I32EnumAttr<"Dimension",
-    "a dimension, either 'x', 'y', or 'z'",
-    [
-      I32EnumAttrCase<"x", 0>,
-      I32EnumAttrCase<"y", 1>,
-      I32EnumAttrCase<"z", 2>
-    ]>{
-  let genSpecializedAttr = 0;
-  let cppNamespace = "::mlir::gpu";
-}
-def GPU_DimensionAttr : EnumAttr<GPU_Dialect, GPU_Dimension, "dim">;
-
 class GPU_IndexOp<string mnemonic, list<Trait> traits = []> :
     GPU_Op<mnemonic, !listconcat(traits, [
         Pure,

mlir/lib/Conversion/GPUCommon/IndexIntrinsicsOpLowering.cpp

@@ -22,53 +22,17 @@ LLVM::ConstantRangeAttr mlir::gpu::index_lowering::getIndexOpRange(
   // 2. Inherent attributes on a surrounding gpu.func
   // 3. Discardable attributes on a surrounding function of any kind
   // The below code handles these in reverse order so that more important
-  // sources overwrite less important ones.
-  DenseI32ArrayAttr funcBounds = nullptr;
-  if (auto funcOp = op->getParentOfType<FunctionOpInterface>()) {
-    switch (indexKind) {
-    case IndexKind::Block: {
-      auto blockHelper =
-          gpu::GPUDialect::KnownBlockSizeAttrHelper(op->getContext());
-      if (blockHelper.isAttrPresent(funcOp))
-        funcBounds = blockHelper.getAttr(funcOp);
-      break;
-    }
-    case IndexKind::Grid: {
-      auto gridHelper =
-          gpu::GPUDialect::KnownGridSizeAttrHelper(op->getContext());
-      if (gridHelper.isAttrPresent(funcOp))
-        funcBounds = gridHelper.getAttr(funcOp);
-      break;
-    }
-    case IndexKind::Cluster: {
-      auto clusterHelper =
-          gpu::GPUDialect::KnownClusterSizeAttrHelper(op->getContext());
-      if (clusterHelper.isAttrPresent(funcOp))
-        funcBounds = clusterHelper.getAttr(funcOp);
-      break;
-    }
-    case IndexKind::Other:
-      break;
-    }
-  }
-  if (auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>()) {
-    switch (indexKind) {
-    case IndexKind::Block:
-      funcBounds = gpuFunc.getKnownBlockSizeAttr();
-      break;
-    case IndexKind::Grid:
-      funcBounds = gpuFunc.getKnownGridSizeAttr();
-      break;
-    case IndexKind::Cluster:
-      funcBounds = gpuFunc.getKnownClusterSizeAttr();
-      break;
-    case IndexKind::Other:
-      break;
-    }
-  }
-  std::optional<uint32_t> upperBound;
-  if (funcBounds)
-    upperBound = funcBounds.asArrayRef()[static_cast<uint32_t>(dim)];
+  // sources overwrite less important ones. As an exception, dimension-size
+  // getters will return exact bounds if known.
+  std::optional<uint32_t> upperBound =
+      getKnownDimensionSizeAround(op, indexKind, dim);
+  // If our upper bound is the maximum possible value, we can't easily construct
+  // the constant range for it.
+  if (upperBound && intrType == IntrType::Dim &&
+      *upperBound < std::numeric_limits<uint32_t>::max())
+    return LLVM::ConstantRangeAttr::get(op->getContext(), bitWidth, *upperBound,
+                                        *upperBound + 1);
+
   if (opUpperBound)
     upperBound = *opUpperBound;
 

mlir/lib/Conversion/GPUCommon/IndexIntrinsicsOpLowering.h

@@ -16,7 +16,9 @@
 namespace mlir {
 namespace gpu {
 namespace index_lowering {
-enum class IndexKind : uint32_t { Other = 0, Block = 1, Grid = 2, Cluster = 3 };
+// Alias so existing call sites don't need updating.
+using IndexKind = gpu::DimensionKind;
+
 enum class IntrType : uint32_t {
   None = 0,
   Id = 1,
@@ -27,7 +29,9 @@ enum class IntrType : uint32_t {
 /// are found. `bitWidth` controls the width of the returned range.
 /// Checks the provided upper_bound from the op (highest priority), inherent
 /// attrs on enclosing `gpu.func`s, and discardable attributes on other
-/// enclosing function ops (lowest priority).
+/// enclosing function ops (lowest priority). However, in the case where
+/// a dimension is known to have a constant value, returns a range indicating
+/// that value.
 LLVM::ConstantRangeAttr getIndexOpRange(Operation *op, gpu::Dimension dim,
                                         std::optional<uint32_t> opUpperBound,
                                         IndexKind indexKind, IntrType intrType,

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -106,19 +106,24 @@ static Value getLaneId(RewriterBase &rewriter, Location loc) {
 static constexpr int64_t kMaxThreadsPerBlockDim = 1024;
 
 /// Emits a call to an OCKL block/grid size function corresponding to
-/// `indexKind` with argument `dim`, querying for upper bounds in the context
-/// surrounding `contextOp` as a fallback for an unknown/unavailable
-/// `opUpperBound`.
-static Value getOcklDim(RewriterBase &rewriter,
-                        gpu::index_lowering::IndexKind indexKind,
-                        gpu::Dimension dim, Operation *contextOp,
-                        std::optional<uint32_t> opUpperBound) {
+/// `indexKind` with argument `dim`, except that if the context around
+/// `contextOp` gives an exact size for that dimension, return that as
+/// an `i64` constant instead.
+static Value getKnownOrOcklDim(RewriterBase &rewriter,
+                               gpu::index_lowering::IndexKind indexKind,
+                               gpu::Dimension dim, Operation *contextOp,
+                               std::optional<uint32_t> opUpperBound) {
   Location loc = contextOp->getLoc();
   MLIRContext *context = contextOp->getContext();
 
   auto i32Ty = IntegerType::get(context, 32);
   auto i64Ty = IntegerType::get(context, 64);
 
+  if (std::optional<uint32_t> knownDim =
+          gpu::getKnownDimensionSizeAround(contextOp, indexKind, dim))
+    return LLVM::ConstantOp::create(rewriter, loc,
+                                    rewriter.getI64IntegerAttr(*knownDim));
+
   int32_t dimParam = static_cast<int32_t>(dim);
 
   StringRef functionName;
@@ -145,14 +150,16 @@ static Value getOcklDim(RewriterBase &rewriter,
   auto callOp =
       LLVM::CallOp::create(rewriter, loc, funcOp, ValueRange{dimConst});
 
-  // Set range attribute on the call result if bounds are available.
-  auto range = gpu::index_lowering::getIndexOpRange(
-      contextOp, dim, opUpperBound, indexKind,
-      gpu::index_lowering::IntrType::Dim, /*bitWidth=*/64);
-  // Fall back to the hardware limit for block dimensions.
-  if (!range && indexKind == gpu::index_lowering::IndexKind::Block)
+  LLVM::ConstantRangeAttr range;
+  if (opUpperBound) {
+    range = LLVM::ConstantRangeAttr::get(
+        context, APInt(64, 1),
+        APInt(64, static_cast<uint64_t>(*opUpperBound) + 1));
+  } else if (indexKind == gpu::index_lowering::IndexKind::Block) {
+    // Set the hardware limit for block ranges as the bounds on block dim calls.
     range = LLVM::ConstantRangeAttr::get(context, APInt(64, 1),
                                          APInt(64, kMaxThreadsPerBlockDim + 1));
+  }
   if (range) {
     callOp.setResAttrsAttr(rewriter.getArrayAttr(rewriter.getDictionaryAttr(
         rewriter.getNamedAttr(LLVM::LLVMDialect::getRangeAttrName(), range))));
@@ -186,8 +193,8 @@ struct GPUDimOpToOcklCall final : ConvertOpToLLVMPattern<OpTy> {
     if (auto bound = op.getUpperBound())
       opUpperBound = static_cast<uint32_t>(bound->getZExtValue());
 
-    Value ocklCall =
-        getOcklDim(rewriter, indexKind, op.getDimension(), op, opUpperBound);
+    Value ocklCall = getKnownOrOcklDim(rewriter, indexKind, op.getDimension(),
+                                       op, opUpperBound);
     Value result = truncOrExtToLLVMType(rewriter, loc, ocklCall,
                                         *this->getTypeConverter());
     rewriter.replaceOp(op, result);
@@ -305,8 +312,8 @@ struct GPUSubgroupIdOpToROCDL : ConvertOpToLLVMPattern<gpu::SubgroupIdOp> {
 
       auto getBlockDim = [&](gpu::Dimension dim) {
         Value dim64 =
-            getOcklDim(rewriter, gpu::index_lowering::IndexKind::Block, dim, op,
-                       std::nullopt);
+            getKnownOrOcklDim(rewriter, gpu::index_lowering::IndexKind::Block,
+                              dim, op, std::nullopt);
         Value dimTrunc =
             LLVM::TruncOp::create(rewriter, loc, int32Type, dim64, flags);
         return dimTrunc;