[mlir][SCFToGPU] LaunchOp propagate optional attributes

[abialas]

Summary

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.

[ftynse]

Consider adding a test.

[abialas]

added simple test

mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

if (isMappedToProcessor(processor)) {
       rewriter.setInsertionPointToStart(loopOp.getBody());
       // Put a sentinel into the worklist so we know when to pop out of the loop
       // body again. We use the launchOp here, as that cannot be part of the
       // bodies instruction.
       worklist.push_back(launchOp.getOperation());
     }
     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()))
     worklist.push_back(&op);
   return success();
 }
 
 /// Lower a `scf.parallel` operation into a corresponding `gpu.launch`

mlir/test/Conversion/SCFToGPU/parallel_loop.mlir

 // CHECK:             gpu.terminator
 // CHECK:           }
 // CHECK:           return
 // CHECK:         }
 // CHECK:       }
 
 // -----
 
+// 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,
                           %arg3 : index,
                           %buf : memref<?x?xf32>,
                           %res : memref<?x?xf32>) {
   %four = constant 4 : index
   // expected-error@+2 {{cannot redefine the bound for processor 1}}