llvm/mlir/include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td

//===- MemRefTransformOps.td - MemRef transformation ops --*- tablegen -*--===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef MEMREF_TRANSFORM_OPS
#define MEMREF_TRANSFORM_OPS

include "mlir/Dialect/Transform/IR/TransformDialect.td"
include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.td"
include "mlir/Dialect/Transform/IR/TransformTypes.td"
include "mlir/Interfaces/SideEffectInterfaces.td"
include "mlir/IR/OpBase.td"

def MemrefToLLVMTypeConverterOp : Op<Transform_Dialect,
    "apply_conversion_patterns.memref.memref_to_llvm_type_converter",
    [DeclareOpInterfaceMethods<TypeConverterBuilderOpInterface,
                               ["getTypeConverter",
                                "getTypeConverterType"]>]> {
  let description = [{
    This operation provides an "LLVMTypeConverter" that lowers memref types to
    LLVM types.

    The type converter can be customized as follows:
    - `use_aligned_alloc`: Use aligned_alloc in place of malloc for heap
      allocations.
    - `index_bitwidth`: Bitwidth of the index type, "0" indicates the size of a
      machine word.
    - `use_generic_functions`: Use generic allocation and deallocation functions
      instead of the classic "malloc", "aligned_alloc" and "free" functions.
    // TODO: the following two options don't really make sense for 
    // memref_to_llvm_type_converter specifically.
    // We should have a single to_llvm_type_converter.
    - `use_bare_ptr_call_conv`: Replace FuncOp's MemRef arguments with bare 
      pointers to the MemRef element types.
    - `data-layout`: String description (LLVM format) of the data layout that is
      expected on the produced module.
  }];

  let arguments = (ins
      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_aligned_alloc,
      DefaultValuedOptionalAttr<I64Attr, "64">:$index_bitwidth,
      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_generic_functions,
      DefaultValuedOptionalAttr<BoolAttr, "false">:$use_bare_ptr_call_conv,
      OptionalAttr<StrAttr>:$data_layout);
  let assemblyFormat = "attr-dict";
}

def ApplyAllocToAllocaOp : Op<Transform_Dialect,
    "apply_patterns.memref.alloc_to_alloca",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface, ["populatePatternsWithState"]>]> {
  let description = [{
    Collects patterns to rewrite scoped dynamic allocation (`alloc`/`dealloc`
    pairs) into automatic allocation (`alloca`) in the same scope, for memrefs
    of static shape.

    The `size_limit` attribute controls the maximum allocated memory (in bytes,
    subject to data layout) for which the pattern applies.
  }];

  let arguments = (ins
    OptionalAttr<I64Attr>:$size_limit);
  let assemblyFormat = "(`size_limit` `(` $size_limit^ `)`)? attr-dict";
}

def ApplyExpandOpsPatternsOp : Op<Transform_Dialect,
    "apply_patterns.memref.expand_ops",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
  let description = [{
    Collects patterns to rewrite ops within the memref dialect.

    - Converts `atomic_rmw` that cannot be lowered to a simple atomic op with
      AtomicRMWOpLowering pattern, e.g. with "minf" or "maxf" attributes, to
      `memref.generic_atomic_rmw` with the expanded code.
    - Converts `memref.reshape` that has a target shape of a statically-known
      size to `memref.reinterpret_cast`.
  }];

  let assemblyFormat = "attr-dict";
}

def ApplyExpandStridedMetadataPatternsOp : Op<Transform_Dialect,
    "apply_patterns.memref.expand_strided_metadata",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
  let description = [{
    Collects patterns for expanding memref operations that modify the metadata
    (sizes, offset, strides) of a memref into easier to analyze constructs.
  }];

  let assemblyFormat = "attr-dict";
}

def ApplyExtractAddressComputationsPatternsOp : Op<Transform_Dialect,
    "apply_patterns.memref.extract_address_computations",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
  let description = [{
    Collects patterns for extracting address computations from operations
    with memory accesses such that these memory accesses use only a base
    pointer.

    For instance,
    ```mlir
    memref.load %base[%off0, ...]
    ```

    Will be rewritten in:
    ```mlir
    %new_base = memref.subview %base[%off0,...][1,...][1,...]
    memref.load %new_base[%c0,...]
    ```
  }];

  let assemblyFormat = "attr-dict";
}

def ApplyFoldMemrefAliasOpsPatternsOp : Op<Transform_Dialect,
    "apply_patterns.memref.fold_memref_alias_ops",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
  let description = [{
    Collects patterns for folding memref aliasing ops (memref.subview) into
    consumer load/store ops (affine.load, memref.load, nvgpu.ldmatrix,
    vector.load, vector.transfer_read, affine.store, memref.store, etc.) and
    other ops (e.g., memref.subview).
  }];

  let assemblyFormat = "attr-dict";
}

def ApplyResolveRankedShapedTypeResultDimsPatternsOp : Op<Transform_Dialect,
    "apply_patterns.memref.resolve_ranked_shaped_type_result_dims",
    [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
  let description = [{
    Collects patterns that resolve `memref.dim` operations with values that are
    defined by operations that implement the `ReifyRankedShapedTypeOpInterface`,
    in terms of shapes of its input operands.
  }];

  let assemblyFormat = "attr-dict";
}

def Transform_MemRefAllocOp : Transform_ConcreteOpType<"memref.alloc">;
def Transform_MemRefAllocaOp : Transform_ConcreteOpType<"memref.alloca">;

def MemRefAllocaToGlobalOp :
  Op<Transform_Dialect, "memref.alloca_to_global",
     [TransformOpInterface,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
  let description = [{
    Inserts a new `memref.global` for each provided `memref.alloca` into the
    nearest symbol table (e.g., a `builtin.module`) and replaces it with a
    `memref.get_global`. This is useful, for example, for allocations that
    should reside in the shared memory of a GPU, which have to be declared as
    globals.

    #### Example

    Consider the following transform op:

    ```mlir
    %get_global, %global =
        transform.memref.alloca_to_global %alloca
          : (!transform.op<"memref.alloca">)
            -> (!transform.any_op, !transform.any_op)
    ```

    and the following input payload:

    ```mlir
    module {
      func.func @func() {
        %alloca = memref.alloca() : memref<2x32xf32>
        // usages of %alloca...
      }
    }
    ```

    then applying the transform op to the payload would result in the following
    output IR:

    ```mlir
    module {
      memref.global "private" @alloc : memref<2x32xf32>
      func.func @func() {
        %alloca = memref.get_global @alloc : memref<2x32xf32>
        // usages of %alloca...
      }
    }
    ```

    #### Return modes

    Succeeds always. The returned handles refer to the `memref.get_global` and
    `memref.global` ops that were inserted by the transformation.
  }];

  let arguments = (ins Transform_MemRefAllocaOp:$alloca);
  let results = (outs TransformHandleTypeInterface:$getGlobal,
                  TransformHandleTypeInterface:$global);

  let assemblyFormat = [{
    $alloca attr-dict `:` functional-type(operands, results)
  }];
}

def MemRefMultiBufferOp : Op<Transform_Dialect, "memref.multibuffer",
    [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
     DeclareOpInterfaceMethods<TransformOpInterface>]> {
  let summary = "Multibuffers an allocation";
  let description = [{
     Transformation to do multi-buffering/array expansion to remove
     dependencies on the temporary allocation between consecutive loop
     iterations. This transform expands the size of an allocation by
     a given multiplicative factor and fixes up any users of the
     multibuffered allocation.
     If skip analysis is not set the transformation will only apply
     if it can prove that there is no data being carried across loop
     iterations.

     #### Return modes

     This operation returns the new allocation if multi-buffering
     succeeds, and failure otherwise.
  }];

  let arguments =
      (ins Transform_MemRefAllocOp:$target,
           ConfinedAttr<I64Attr, [IntPositive]>:$factor,
           UnitAttr:$skip_analysis);

  let results = (outs TransformHandleTypeInterface:$transformed);

  let assemblyFormat =
    "$target attr-dict `:` functional-type(operands, results)";
}

def MemRefEraseDeadAllocAndStoresOp
    : Op<Transform_Dialect, "memref.erase_dead_alloc_and_stores", [
      TransformEachOpTrait, TransformOpInterface,
      DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
      ReportTrackingListenerFailuresOpTrait
    ]> {
  let description = [{
    This applies memory optimization on memref. In particular it does store to
    load forwarding, dead store elimination and dead alloc elimination.

    #### Return modes

    This operation applies a set of memory optimization on the whole region of
    the operand.

    The transformation does not consume the target handle. It modifies the
    payload. Dead allocations, loads and stores are silently dropped from all
    mappings.
  }];

  let arguments = (ins TransformHandleTypeInterface:$target);
  let results = (outs);

  let assemblyFormat = "$target attr-dict `:` functional-type($target, results)";

  let skipDefaultBuilders = 1;
  let builders = [
    OpBuilder<(ins "Value":$target)>
  ];
  let extraClassDeclaration = [{
    ::mlir::DiagnosedSilenceableFailure applyToOne(
        ::mlir::transform::TransformRewriter &rewriter,
        ::mlir::Operation *target,
        ::mlir::transform::ApplyToEachResultList &results,
        ::mlir::transform::TransformState &state);
  }];
}

def MemRefMakeLoopIndependentOp
    : Op<Transform_Dialect, "memref.make_loop_independent",
         [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
          TransformOpInterface, TransformEachOpTrait]> {
  let description = [{
    Rewrite the targeted ops such that their index-typed operands no longer
    depend on any loop induction variable of the `num_loop` enclosing `scf.for`
    loops. I.e., compute an upper bound that is independent of any such loop IV
    for every tensor dimension. The transformed op could then be hoisted from
    the `num_loop` enclosing loops. To preserve the original semantics, place a
    `memref.subview` inside the loop.

    Currently supported operations are:
    - memref.alloca: Replaced with a new memref.alloca with upper bound sizes,
      followed by a memref.subview.

    #### Return modes

    This operation fails if at least one induction variable could not be
    eliminated. In case the targeted op is already independent of induction
    variables, this transform succeeds and returns the unmodified target op.

    Otherwise, the returned handle points to a subset of the produced ops:
    - memref.alloca: The returned handle points to the memref.subview op.

    This transform op consumes the target handle and produces a result handle.
  }];

  let arguments = (ins TransformHandleTypeInterface:$target, I64Attr:$num_loops);
  let results = (outs TransformHandleTypeInterface:$transformed);
  let assemblyFormat =
      "$target attr-dict `:` functional-type($target, $transformed)";

  let extraClassDeclaration = [{
    ::mlir::DiagnosedSilenceableFailure applyToOne(
        ::mlir::transform::TransformRewriter &rewriter,
        ::mlir::Operation *target,
        ::mlir::transform::ApplyToEachResultList &results,
        ::mlir::transform::TransformState &state);
  }];
}

#endif // MEMREF_TRANSFORM_OPS