[mlir][gpu] sinkOperationsIntoLaunchOp: Add user hook for isSinkingBeneficiary

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/Passes.h

30 ↗

(On Diff #408183)

I'm concerned about making any pass non hermetic: that mean adding any C++ injection that won't be serialized in the textual pipeline.

I'd like to keep the pass simple and hermetic, and extract whatever logic as utilities so that one can customize in their own pass instead.

[Hardcode84]

revert pass changes, add unittest

[Hardcode84]

fix licence header

[herhut]

Thank you. The default is fairly arbitrary (well, it's good enough for some index computations) and we really should have some cost model here. +1 to making it configurable at least.

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

128

It would be nice to maybe split the pass into two, one that does sinking and one that does outlining. Then we could remove the default here completely and add it to the sinking pass instead. Actual users of this API would need to implement their own cost model.

[Hardcode84]

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

128

I can make isSinkingBeneficiary param non-optional and provide isSinkingBeneficiaryDefault explicitly in GpuKernelOutliningPass but I would prefer to do pass splitting in separate review.

[Hardcode84]

make isSinkingBeneficiary param non-optional

[mehdi_amini]

mlir/unittests/Dialect/GPU/KernelOutliningTest.cpp

57 ↗

(On Diff #408490)

Please keep the testing implemented with passes, this is how MLIR is tested in general.

[Hardcode84]

mlir/unittests/Dialect/GPU/KernelOutliningTest.cpp

57 ↗

(On Diff #408490)

I that case I can just drop these test and assume it is tested by existing GpuKernelOutliningPass tests.

[herhut]

Thank you. Just a naming nit.

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

128

Sounds good. Maybe also remove the 'Default' from the name then. We could even call it isLikelyAnIndexComputation or something because that is what the function really does.

mlir/unittests/Dialect/GPU/KernelOutliningTest.cpp

57 ↗

(On Diff #408490)

If we do not have a default implementation, then this test is not needed indeed. The other test would already test that providing an implementation works and we do not have to test the overriding.

[Hardcode84]

rebase, remove unittest, rename default func

mlir/include/mlir/Dialect/GPU/Utils.h

 ///   `launchOp` and defined above it. These captured values are appended to the
 ///   `operands` list.
 gpu::GPUFuncOp outlineKernelFunc(gpu::LaunchOp launchOp, StringRef kernelFnName,
                                  SmallVectorImpl<Value> &operands);
 
 /// Sink operations into the `launchOp` to reduce the number of values that are
 /// used within the region of the operation, but defined outside of the
 /// region.
-LogicalResult sinkOperationsIntoLaunchOp(gpu::LaunchOp launchOp);
+LogicalResult sinkOperationsIntoLaunchOp(
+    gpu::LaunchOp launchOp,
+    llvm::function_ref<bool(Operation *)> isSinkingBeneficiary);
 
 } // namespace mlir
 #endif // MLIR_DIALECT_GPU_UTILS_H_

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

static void injectGpuIndexOperations(Location loc, Region &launchFuncOpBody,
   // operations. Iterate backwards since args are erased and indices change.
   for (const auto &indexOp : enumerate(indexOps))
     map.map(firstBlock.getArgument(indexOp.index()), indexOp.value());
 }
 
 /// Identifies operations that are beneficial to sink into kernels. These
 /// operations may not have side-effects, as otherwise sinking (and hence
 /// duplicating them) is not legal.
-static bool isSinkingBeneficiary(Operation *op) {
+static bool isLikelyAnIndexComputatio(Operation *op) {
   return isa<arith::ConstantOp, ConstantOp, memref::DimOp, arith::SelectOp,
              arith::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 `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 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,
-                      const SetVector<Value> &existingDependencies,
-                      SetVector<Operation *> &beneficiaryOps,
-                      llvm::SmallPtrSetImpl<Value> &availableValues) {
+static bool extractBeneficiaryOps(
+    Operation *op, const SetVector<Value> &existingDependencies,
+    llvm::SmallPtrSetImpl<Value> &availableValues,
+    llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
   if (beneficiaryOps.count(op))
     return true;
 
   if (!isSinkingBeneficiary(op))
     return false;
 
   for (Value operand : op->getOperands()) {
     // It is already visible in the kernel, keep going.
     if (availableValues.count(operand))
       continue;
     // Else check whether it can be made available via sinking or already is a
     // dependency.
     Operation *definingOp = operand.getDefiningOp();
-    if ((!definingOp ||
-         !extractBeneficiaryOps(definingOp, existingDependencies,
-                                beneficiaryOps, availableValues)) &&
+    if ((!definingOp || !extractBeneficiaryOps(definingOp, existingDependencies,
+                                               beneficiaryOps, availableValues,
+                                               isSinkingBeneficiary)) &&
         !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) {
+LogicalResult mlir::sinkOperationsIntoLaunchOp(
+    gpu::LaunchOp launchOp,
+    llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
+  assert(isSinkingBeneficiary);
   Region &launchOpBody = launchOp.body();
 
   // Identify uses from values defined outside of the scope of the launch
   // operation.
   SetVector<Value> sinkCandidates;
   getUsedValuesDefinedAbove(launchOpBody, sinkCandidates);
 
   SetVector<Operation *> toBeSunk;
   llvm::SmallPtrSet<Value, 4> availableValues;
   for (Value operand : sinkCandidates) {
     Operation *operandOp = operand.getDefiningOp();
     if (!operandOp)
       continue;
-    extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues);
+    extractBeneficiaryOps(operandOp, sinkCandidates, toBeSunk, availableValues,
+                          isSinkingBeneficiary);
   }
 
   // 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.
     for (auto pair : llvm::zip(op->getResults(), clonedOp->getResults()))
for (auto func : getOperation().getOps<FuncOp>()) {
       // Insert just after the function.
       Block::iterator insertPt(func->getNextNode());
       auto funcWalkResult = func.walk([&](gpu::LaunchOp op) {
         SetVector<Value> operands;
         std::string kernelFnName =
             Twine(op->getParentOfType<FuncOp>().getName(), "_kernel").str();
 
         // Pull in instructions that can be sunk
-        if (failed(sinkOperationsIntoLaunchOp(op)))
+        if (failed(sinkOperationsIntoLaunchOp(op, isLikelyAnIndexComputatio)))
           return WalkResult::interrupt();
         gpu::GPUFuncOp outlinedFunc =
             outlineKernelFuncImpl(op, kernelFnName, operands);
 
         // Create nested module and insert outlinedFunc. The module will
         // originally get the same name as the function, but may be renamed on
         // insertion into the parent module.
         auto kernelModule = createKernelModule(outlinedFunc, symbolTable);