[Mlir-commits] [mlir] 396e7f4 - [mlir][SCFToGPU] LaunchOp propagate optional attributes

[Artur Bialas]

Author: Artur Bialas
Date: 2020-09-25T09:21:16+02:00
New Revision: 396e7f454893e24969bb989fe89aa028e2ea1693

URL: https://github.com/llvm/llvm-project/commit/396e7f454893e24969bb989fe89aa028e2ea1693
DIFF: https://github.com/llvm/llvm-project/commit/396e7f454893e24969bb989fe89aa028e2ea1693.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
    mlir/test/Conversion/SCFToGPU/parallel_loop.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp b/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
index b1d5a854de80..21a698628737 100644
--- a/mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp
+++ b/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()))

diff  --git a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir b/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir
index b4131a3fda0e..5e3e2dc5aa82 100644
--- a/mlir/test/Conversion/SCFToGPU/parallel_loop.mlir
+++ b/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,