Linalg copy instead of memref copy when non identity layout

[rmoyard]

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 with linalg.copy, it lowers to scf and standard dialects and does not call a runtime function.

Benefits: More at compilation time, and compatibility with Enzyme.

[rmoyard]

[sc-67739]

[erick-xanadu]

@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.

[erick-xanadu]

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.

[codecov[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 97.93%. Comparing base (498210c) to head (2af87ac).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #917 +/- ## ========================================== - Coverage 99.90% 97.93% -1.97% ========================================== Files 21 73 +52 Lines 4154 10338 +6184 Branches 200 1170 +970 ========================================== + Hits 4150 10125 +5975 - Misses 4 170 +166 - Partials 0 43 +43 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.