/// 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);
}
/// 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,
+ SetVector<Operation *> &beneficiaryOps,
+ llvm::SmallPtrSetImpl<Value> &availableValues,
+ llvm::function_ref<bool(Operation *)> isSinkingBeneficiary) {
if (beneficiaryOps.count(op))
return true;
// 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;
}
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 *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.
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);