diff --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp --- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp +++ b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp @@ -517,6 +517,16 @@ } cloningMap.map(iv, newIndex); } + + // Propagate custom user defined optional attributes, that can be used at + // later stage, such as extension data for GPU kernel dispatch + for (const auto &namedAttr : parallelOp.getAttrs()) { + if (namedAttr.first == gpu::getMappingAttrName() || + namedAttr.first == ParallelOp::getOperandSegmentSizeAttr()) + continue; + launchOp.setAttr(namedAttr.first, namedAttr.second); + } + Block *body = parallelOp.getBody(); worklist.reserve(worklist.size() + body->getOperations().size()); for (Operation &op : llvm::reverse(body->without_terminator())) diff --git a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir b/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir --- a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir +++ b/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir @@ -304,6 +304,19 @@ // ----- +// Optional attribute lowering test + +func @parallel_loop_optional_attr() { + %c0 = constant 0 : index + %c1 = constant 1 : index + scf.parallel (%i0) = (%c0) to (%c1) step (%c1) { + } { mapping = [{processor = 0, map = affine_map<(d0) -> (d0)>, bound = affine_map<(d0) -> (d0)>}], optional_attr = 1 } + // CHECK: optional_attr = 1 + return +} + +// ----- + // Mapping to the same processor twice. func @parallel_double_map(%arg0 : index, %arg1 : index, %arg2 : index,