gpu.launch_func: add accessors to grid configuration operands

Add accessor functions that return `gpu::KernelDim3` containing the respective
operands for grid and block size accepted by `gpu.launch_func`.  Use the same
signature as for `gpu.launch`.

PiperOrigin-RevId: 254942674

diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h
index b06e6a7..d48cc1a 100644 (file)
--- a/mlir/include/mlir/GPU/GPUDialect.h
+++ b/mlir/include/mlir/GPU/GPUDialect.h
@@ -133,6 +133,11 @@ public:
   /// The i-th operand passed to the kernel function.
   Value *getKernelOperand(unsigned i);
 
+  /// Get the SSA values passed as operands to specify the grid size.
+  KernelDim3 getGridSizeOperandValues();
+  /// Get the SSA values passed as operands to specify the block size.
+  KernelDim3 getBlockSizeOperandValues();
+
   LogicalResult verify();
 
   static StringRef getOperationName() { return "gpu.launch_func"; }

diff --git a/mlir/lib/GPU/IR/GPUDialect.cpp b/mlir/lib/GPU/IR/GPUDialect.cpp
index 272ffce..c00b492 100644 (file)
--- a/mlir/lib/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/GPU/IR/GPUDialect.cpp
@@ -344,6 +344,14 @@ Value *LaunchFuncOp::getKernelOperand(unsigned i) {
   return getOperation()->getOperand(i + kNumConfigOperands);
 }
 
+KernelDim3 LaunchFuncOp::getGridSizeOperandValues() {
+  return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
+}
+
+KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() {
+  return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
+}
+
 LogicalResult LaunchFuncOp::verify() {
   auto kernelAttr = this->getAttr(getKernelAttrName());
   if (!kernelAttr) {