[mlir][SCFToGPU] LaunchOp propagate optional attributes

Allow propagating optional user defined attributes during SCF to GPU conversion. Gives opportunity to use user defined attributes in the further lowering. For example setting subgroup size, or other options for GPU dispatch. This does not break backward compatibility and does not require new attributes, just allow passing optional ones.

Differential Revision: https://reviews.llvm.org/D88203

mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

@@ -517,6 +517,16 @@ static LogicalResult processParallelLoop(
     }
     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()))

mlir/test/Conversion/SCFToGPU/parallel_loop.mlir

@@ -304,6 +304,19 @@ module {
 
 // -----
 
+// 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,