Closed rmoyard closed 1 month ago
[sc-67739]
@rmoyard , in the discord channel someone asked a very similar question. There was this answer:
There is an option in the bufferization passes where you can customize which copy op is used, and there you can tell it to use linalg.copy
DiagnosedSilenceableFailure
transform::OneShotBufferizeOp::apply(transform::TransformRewriter &rewriter,
TransformResults &transformResults,
TransformState &state) {
OneShotBufferizationOptions options;
options.allowReturnAllocsFromLoops = getAllowReturnAllocsFromLoops();
options.allowUnknownOps = getAllowUnknownOps();
options.bufferizeFunctionBoundaries = getBufferizeFunctionBoundaries();
options.dumpAliasSets = getDumpAliasSets();
options.testAnalysisOnly = getTestAnalysisOnly();
options.printConflicts = getPrintConflicts();
if (getFunctionBoundaryTypeConversion().has_value())
options.setFunctionBoundaryTypeConversion(
*getFunctionBoundaryTypeConversion());
if (getMemcpyOp() == "memref.copy") {
options.memCpyFn = [](OpBuilder &b, Location loc, Value from, Value to) {
b.create<memref::CopyOp>(loc, from, to);
return success();
};
} else if (getMemcpyOp() == "linalg.copy") {
options.memCpyFn = [](OpBuilder &b, Location loc, Value from, Value to) {
b.create<linalg::CopyOp>(loc, from, to);
return success();
};
} else {
llvm_unreachable("invalid copy op");
}
It looks like in order to do this from the one-shot bufferizer we might just need to extend a class and implement getMemcpyOp
to return "linalg.copy", but I am not 100% sure.
Nevermind what I said, it is using the transform dialect to schedule bufferization. Still an interesting approach. Pinging @paul0403 since he is working on the transform dialect right now.
def OneShotBufferizeOp
: Op<Transform_Dialect, "bufferization.one_shot_bufferize",
[FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
DeclareOpInterfaceMethods<TransformOpInterface>]> {
let description = [{
Indicates that the given `target` op should be bufferized with One-Shot
Bufferize. The bufferization can be configured with various attributes that
corresponding to options in `BufferizationOptions` and the
`one-shot-bufferize` pass. More information can be found in the pass
documentation.
The targeted ops must be modules or functions. This is because there is
always a single, bufferized replacement op for such targets.
Note: Only ops that implement `BufferizableOpInterface` are bufferized. All
other ops are ignored if `allow_unknown_ops`. If `allow_unknown_ops` is
unset, this transform fails when an unknown/non-bufferizable op is found.
Many ops implement `BufferizableOpInterface` via an external model. These
external models must be registered when applying this transform op;
otherwise, said ops would be considered non-bufferizable.
#### Return modes
This operation consumes the `target` handle and produces the `transformed`
handle.
}];
let arguments = (
ins TransformHandleTypeInterface:$target,
OptionalAttr<LayoutMapOption>:$function_boundary_type_conversion,
DefaultValuedAttr<BoolAttr, "false">:$allow_return_allocs_from_loops,
DefaultValuedAttr<BoolAttr, "false">:$allow_unknown_ops,
DefaultValuedAttr<BoolAttr, "false">:$bufferize_function_boundaries,
DefaultValuedAttr<BoolAttr, "false">:$dump_alias_sets,
DefaultValuedAttr<BoolAttr, "false">:$test_analysis_only,
DefaultValuedAttr<BoolAttr, "false">:$print_conflicts,
DefaultValuedAttr<StrAttr, "\"memref.copy\"">:$memcpy_op);
let results = (outs TransformHandleTypeInterface:$transformed);
let hasVerifier = 1;
let assemblyFormat = [{
(`layout` `{` $function_boundary_type_conversion^ `}`)?
$target attr-dict `:` functional-type($target, results)
}];
}
But feel free to continue with your approach.
All modified and coverable lines are covered by tests :white_check_mark:
Project coverage is 97.93%. Comparing base (
498210c
) to head (2af87ac
).
:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.
Context:
memref.copy
operation lowers to a runtime function implemented in the mlir execution engine when the layout is not the identity. This causes issues with Enzyme.Description of the Change: When the layout is not the identity we replace
memref.copy
withlinalg.copy
, it lowers to scf and standard dialects and does not call a runtime function.Benefits: More at compilation time, and compatibility with Enzyme.