[mlir][gpu] Fix bug in kernel outlining

[herhut]

Summary

The updated version of kernel outlining did not handle cases correctly
where an operand of a candidate for sinking itself was defined by an operation
that is a sinking candidate. In such cases, it could happen that sunk
operations were inserted in the wrong order, breaking ssa properties.

[bondhugula]

mlir/test/Dialect/GPU/outlining.mlir

197–201

Nit: consider indenting the block.

[herhut]

Formatting

mlir/lib/Dialect/GPU/Transforms/KernelOutlining.cpp

 static bool isSinkingBeneficiary(Operation *op) {
   return isa<ConstantOp, DimOp, SelectOp, CmpIOp>(op);
 }
 
 /// For a given operation `op`, computes whether it is beneficial to sink the
 /// operation into the kernel. An operation can be sunk if doing so does not
 /// introduce new kernel arguments. Whether a value is already available in the
 /// kernel (and hence does not introduce new arguments) is checked by
-/// querying `availableValues`.
+/// querying `existingDependencies` and `availableValues`.
 /// If an operand is not yet available, we recursively check whether it can be
 /// made available by siking its defining op.
 /// Operations that are indentified for sinking are added to `beneficiaryOps` in
-/// the order the should appear in the kernel. Furthermore, `availableValues` is
-/// updated with results that will be available after sinking the identified
+/// the order they should appear in the kernel. Furthermore, `availableValues`
+/// is updated with results that will be available after sinking the identified
 /// ops.
-static bool extractBeneficiaryOps(Operation *op,
-                                  llvm::SetVector<Operation *> &beneficiaryOps,
-                                  llvm::SetVector<Value> &availableValues) {
+static bool
+extractBeneficiaryOps(Operation *op,
+                      llvm::SetVector<Value> existingDependencies,
+                      llvm::SmallPtrSetImpl<Value> &availableValues) {
   if (beneficiaryOps.count(op))
     return true;
 
   if (!isSinkingBeneficiary(op))
     return false;
 
   for (Value operand : op->getOperands()) {
     // It is already visisble in the kernel, keep going.
     if (availableValues.count(operand))
       continue;
-    // Else check whether it can be made available via sinking.
+    // Else check whether it can be made available via sinking or already is a
+    // dependency.
     Operation *definingOp = operand.getDefiningOp();
-    if (!definingOp ||
-        !extractBeneficiaryOps(definingOp, beneficiaryOps, availableValues))
+    if ((!definingOp ||
+         !extractBeneficiaryOps(definingOp, existingDependencies,
+                                beneficiaryOps, availableValues)) &&
+        !existingDependencies.count(operand))
       return false;
   }
   // We will sink the operation, mark its results as now available.
   beneficiaryOps.insert(op);
   for (Value result : op->getResults())
     availableValues.insert(result);
   return true;
 }
 
 LogicalResult mlir::sinkOperationsIntoLaunchOp(gpu::LaunchOp launchOp) {
   Region &launchOpBody = launchOp.body();
 
   // Identify uses from values defined outside of the scope of the launch
   // operation.
   llvm::SetVector<Value> sinkCandidates;
   getUsedValuesDefinedAbove(launchOpBody, sinkCandidates);
 
-  SmallVector<Value, 4> worklist(sinkCandidates.begin(), sinkCandidates.end());
   llvm::SetVector<Operation *> toBeSunk;
-  for (Value operand : worklist) {
+  llvm::SmallPtrSet<Value, 4> availableValues;
+  for (Value operand : sinkCandidates) {
     Operation *operandOp = operand.getDefiningOp();
     if (!operandOp)
       continue;
-    extractBeneficiaryOps(operandOp, toBeSunk, sinkCandidates);
+    extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues);
   }
 
   // Insert operations so that the defs get cloned before uses.
   BlockAndValueMapping map;
   OpBuilder builder(launchOpBody);
   for (Operation *op : toBeSunk) {
     Operation *clonedOp = builder.clone(*op, map);
     // Only replace uses within the launch op.

mlir/test/Dialect/GPU/outlining.mlir

 
 // CHECK-LABEL: func @extra_constants_noarg_kernel
 // CHECK-SAME: %[[KARG0:.*]]: memref<?xf32>, %[[KARG1:.*]]: index
 // CHECK: %[[KCST:.*]] = constant 2
 // CHECK: "use"(%[[KCST]], %[[KARG0]], %[[KARG1]])
 
 // -----
 
+// CHECK-LABEL: @multiple_uses
 func @multiple_uses(%arg0 : memref<?xf32>) {
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   // CHECK: gpu.func {{.*}} {
-  // CHECK: %[[C2:.*]] = constant 2 : index
-  // CHECK: "use1"(%[[C2]], %[[C2]])
-  // CHECK: "use2"(%[[C2]])
-  // CHECK: gpu.return
+  // CHECK:   %[[C2:.*]] = constant 2 : index
+  // CHECK:   "use1"(%[[C2]], %[[C2]])
+  // CHECK:   "use2"(%[[C2]])
+  // CHECK:   gpu.return
   // CHECK: }
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,
                                        %grid_z = %c1)
              threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,
                                         %block_z = %c1) {
     "use1"(%c2, %c2) : (index, index) -> ()
     "use2"(%c2) : (index) -> ()
     gpu.terminator
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: @multiple_uses2
+func @multiple_uses2(%arg0 : memref<*xf32>) {
+  %c1 = constant 1 : index
+  %c2 = constant 2 : index
+  %d = dim %arg0, %c2 : memref<*xf32>
+  // CHECK: gpu.func {{.*}} {
+  // CHECK:   %[[C2:.*]] = constant 2 : index
+  // CHECK:   %[[D:.*]] = dim %[[ARG:.*]], %[[C2]]
+  // CHECK:   "use1"(%[[D]])
+  // CHECK:   "use2"(%[[C2]], %[[C2]])
+  // CHECK:   "use3"(%[[ARG]])
+  // CHECK:   gpu.return
+  // CHECK: }
+  gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %c1, %grid_y = %c1,
+                                       %grid_z = %c1)
+             threads(%tx, %ty, %tz) in (%block_x = %c1, %block_y = %c1,
+                                        %block_z = %c1) {
+    "use1"(%d) : (index) -> ()
+    "use2"(%c2, %c2) : (index, index) -> ()
+    "use3"(%arg0) : (memref<*xf32>) -> ()
+    gpu.terminator
+  }
+  return
+}
+
+// -----
+
 llvm.mlir.global internal @global(42 : i64) : !llvm.i64
 
 //CHECK-LABEL: @function_call
 func @function_call(%arg0 : memref<?xf32>) {
   %cst = constant 8 : index
   gpu.launch blocks(%bx, %by, %bz) in (%grid_x = %cst, %grid_y = %cst,
                                        %grid_z = %cst)
              threads(%tx, %ty, %tz) in (%block_x = %cst, %block_y = %cst,