include/mlir/Dialect/MemRef/TransformOps/MemRefTransformOps.td - llvm-project/mlir - Git at Google

 //===- 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
	//===- 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