//===- NVGPUTransformOps.td - NVGPU transform 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 NVGPU_TRANSFORM_OPS
#define NVGPU_TRANSFORM_OPS
include "mlir/Dialect/Transform/IR/TransformAttrs.td"
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"
//===----------------------------------------------------------------------===//
// Apply...ConversionPatternsOp
//===----------------------------------------------------------------------===//
def ApplyNVGPUToNVVMConversionPatternsOp : Op<Transform_Dialect,
"apply_conversion_patterns.nvgpu.nvgpu_to_nvvm",
[DeclareOpInterfaceMethods<ConversionPatternDescriptorOpInterface,
["verifyTypeConverter"]>]> {
let description = [{
Collects patterns that convert NVGPU dialect ops to NVVM dialect ops. These
patterns require an "LLVMTypeConverter".
}];
let assemblyFormat = "attr-dict";
}
//===----------------------------------------------------------------------===//
// CreateAsyncGroupsOp
//===----------------------------------------------------------------------===//
def CreateAsyncGroupsOp :
Op<Transform_Dialect, "nvgpu.create_async_groups",
[DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
TransformEachOpTrait,
TransformOpInterface,
ReportTrackingListenerFailuresOpTrait]> {
let description = [{
Look for global to shared memory copies within the targeted op in the form
of vector transfer ops and convert them to async copies when possible.
Consecutive copies are put into the same group. A "wait" operation is
inserted right at the of end the group.
`bypass_l1` specifies whether `bypassL1` attributes should be added to
the async copies. `bypass_l1` is a compiler hint: only 16 byte transfers
can bypass the L1 cache, so this attribute is not set for any other transfer
sizes.
#### Return modes
This op consumes the `target` handle and produces the `result` handle, which
is mapped to the same payload operations as the `target` handle. The op
modifies the payload.
}];
let arguments = (ins TransformHandleTypeInterface:$target,
UnitAttr:$bypass_l1);
let results = (outs TransformHandleTypeInterface:$result);
let assemblyFormat = [{
$target attr-dict `:` functional-type(operands, results)
}];
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::transform::TransformRewriter &rewriter,
::mlir::Operation *target,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
//===----------------------------------------------------------------------===//
// PipelineSharedMemoryCopiesOp
//===----------------------------------------------------------------------===//
def PipelineSharedMemoryCopiesOp :
Op<Transform_Dialect, "nvgpu.pipeline_shared_memory_copies",
[FunctionalStyleTransformOpTrait,
MemoryEffectsOpInterface,
TransformEachOpTrait,
TransformOpInterface,
ReportTrackingListenerFailuresOpTrait]> {
let summary =
"Applies software pipelining to a given loop with shared memory copies";
let description = [{
Applies software pipelining to a given scf.for loop. The pipelining
strategy will look for a load into shared memory and pipeline it to overlap
it with the rest of the loop.
NOTE: It is user responsibility to ensure that there are no dependency
between `depth` iterations of the loop by using multi-buffering. It is
also user responsibility to ensure a sufficient amount of shared memory
is allocated to cover eventual writes by `depth-1` speculative
iterations.
`depth` will indicate how many stages the software pipeline should have.
`peel_epilogue` allows to force the epilogue to be peeled out instead of
potentially using predicated operations for the epilogue phase.
#### Return modes
Consumes the operand handle and produces a result handle pointing to the
loop, which may or may not have been pipelined. Produces a definite failure
if the loop pipeliner mutated the IR before failing to pipeline, in
particular if `peel_epilogue` is not set and the loop body doesn't support
predication. If failure propagation mode is set to "propagate", produces a
silenceable failure when pipelining preconditions, e.g., loop bound being
static, are not met or when the loop wasn't pipelined because due to the
lack of loads into shared memory. If the failure propagation mode is set
to "suppress" (default), succeeds in these case and associates the result
handle with the original loop.
TODO: the shared memory part and behavior specific to NVGPU should be
made orthogonal to pipelining so that `transform.loop.pipeline` becomes
usable here.
}];
let arguments = (ins TransformHandleTypeInterface:$for_op,
I64Attr:$depth,
UnitAttr:$peel_epilogue,
DefaultValuedAttr<FailurePropagationMode,
"::mlir::transform::FailurePropagationMode::Suppress">
:$failure_propagation_mode);
let results = (outs TransformHandleTypeInterface:$result);
let assemblyFormat = [{
`failures` `(` $failure_propagation_mode `)`
$for_op
attr-dict
`:` functional-type(operands, results)
}];
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::transform::TransformRewriter &rewriter,
::mlir::scf::ForOp forOp,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
//===----------------------------------------------------------------------===//
// RewriteMatmulAsMmaSyncOp
//===----------------------------------------------------------------------===//
def RewriteMatmulAsMmaSyncOp :
Op<Transform_Dialect, "nvgpu.rewrite_matmul_as_mma_sync",
[FunctionalStyleTransformOpTrait,
MemoryEffectsOpInterface,
TransformEachOpTrait,
TransformOpInterface,
ReportTrackingListenerFailuresOpTrait]> {
let description = [{
Rewrite a matmul operation on memref to an mma.sync operation on vectors.
Memory copies with the required access patterns are automatically inserted.
Operations that do not have a 1-1 mapping to mma.sync operations are left
unchanged.
}];
let arguments = (ins TransformHandleTypeInterface:$target);
let results = (outs);
let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure applyToOne(
::mlir::transform::TransformRewriter &rewriter,
::mlir::linalg::LinalgOp linalgOp,
::mlir::transform::ApplyToEachResultList &results,
::mlir::transform::TransformState &state);
}];
}
//===----------------------------------------------------------------------===//
// RewriteCopyAsTmaOp
//===----------------------------------------------------------------------===//
def RewriteCopyAsTmaOp :
Op<Transform_Dialect, "nvgpu.rewrite_copy_as_tma",
[FunctionalStyleTransformOpTrait,
MemoryEffectsOpInterface,
TransformEachOpTrait,
TransformOpInterface,
ReportTrackingListenerFailuresOpTrait]> {
let description = [{
Rewrite a copy operation on memref to tma operations that transit through
shared memory.
}];
let arguments = (ins TransformHandleTypeInterface:$target);
let results = (outs);
let assemblyFormat = "$target attr-dict `:` functional-type(operands, results) ";
let extraClassDeclaration = [{
::mlir::DiagnosedSilenceableFailure apply(
::mlir::transform::TransformRewriter &rewriter,
::mlir::transform::TransformResults &transformResults,
::mlir::transform::TransformState &state);
}];
}
#endif // NVGPU_TRANSFORM_OPS
Codebase Browser
llvm