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
This commit is contained in:
Alex Zinenko 2019-06-25 04:21:10 -07:00 committed by A. Unique TensorFlower
parent 14a10a818b
commit 43c0cf109f
2 changed files with 13 additions and 0 deletions

View File

@ -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"; }

View File

@ -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) {