From 43c0cf109ffcec12bb59d1cb5452c66ab074983a Mon Sep 17 00:00:00 2001 From: Alex Zinenko Date: Tue, 25 Jun 2019 04:21:10 -0700 Subject: [PATCH] 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 --- mlir/include/mlir/GPU/GPUDialect.h | 5 +++++ mlir/lib/GPU/IR/GPUDialect.cpp | 8 ++++++++ 2 files changed, 13 insertions(+) diff --git a/mlir/include/mlir/GPU/GPUDialect.h b/mlir/include/mlir/GPU/GPUDialect.h index b06e6a7..d48cc1a 100644 --- 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 --- 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) { -- 2.7.4