1//===- GPUTransformOps.h - GPU transform ops --------------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef MLIR_DIALECT_GPU_TRANSFORMOPS_GPUTRANSFORMOPS_H
10#define MLIR_DIALECT_GPU_TRANSFORMOPS_GPUTRANSFORMOPS_H
11
12#include "mlir/Dialect/SCF/IR/SCF.h"
13#include "mlir/Dialect/Transform/Interfaces/TransformInterfaces.h"
14#include "mlir/IR/OpImplementation.h"
15#include "mlir/IR/PatternMatch.h"
16
17namespace mlir {
18namespace gpu {
19class GpuOp;
20} // namespace gpu
21} // namespace mlir
22
23//===----------------------------------------------------------------------===//
24// GPU Transform Operations
25//===----------------------------------------------------------------------===//
26
27#define GET_OP_CLASSES
28#include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h.inc"
29
30namespace mlir {
31class DialectRegistry;
32namespace transform {
33namespace gpu {
34struct GpuIdBuilder;
35
36/// Map the top level `scf.forall` op to GPU blocks.
37/// Mapping is one-to-one and the induction variables of `scf.forall` are
38/// rewritten to gpu.block_id according to the thread_dim_mapping attribute.
39///
40/// Dynamic, `scf.forall` trip counts are currently not supported.
41/// Dynamic `gridDims` are currently not supported.
42DiagnosedSilenceableFailure
43mapForallToBlocksImpl(RewriterBase &rewriter, TransformOpInterface transformOp,
44 scf::ForallOp forallOp,
45 SmallVectorImpl<int64_t> &gridDims,
46 const GpuIdBuilder &gpuIdBuilder);
47
48/// Search `scf.forall` ops nested under `target` and map each such op to an
49/// explicit GPU implementation along `blockDims`.
50/// The mapping is one-to-one and the induction variables of `scf.forall` are
51/// rewritten to gpuIdBuilder.idBuilder according to the
52/// gpuIdBuilder.mappingAttributes attribute.
53///
54/// Dynamic, `scf.forall` trip counts are currently not supported.
55/// Dynamic `blockDims` sizes are currently not supported.
56/// `blockDims` is expected to be of size 3.
57DiagnosedSilenceableFailure
58mapOneForallToThreadsImpl(RewriterBase &rewriter,
59 std::optional<TransformOpInterface> transformOp,
60 scf::ForallOp forallOp, ArrayRef<int64_t> blockDims,
61 int64_t warpSize, bool syncAfterDistribute);
62
63/// Search `scf.forall` ops nested under `target` and map each such op to an
64/// explicit GPU implementation along `blockDims`.
65/// The mapping is one-to-one and the induction variables of `scf.forall` are
66/// rewritten to appropriate ids according to the mapping attribute.
67///
68/// Dynamic, `scf.forall` trip counts are currently not supported.
69/// Dynamic `blockDims` or `newBasis` entries are currently not
70/// supported. `blockDims` is expected to be of size 3.
71///
72/// The insertion point of the `rewriter` is expected to be set at the
73/// beginning of the `target` body block and dominate all other blocks.
74DiagnosedSilenceableFailure
75mapNestedForallToThreadsImpl(RewriterBase &rewriter,
76 std::optional<TransformOpInterface> transformOp,
77 Operation *target, ArrayRef<int64_t> blockDims,
78 int64_t warpSize, bool syncAfterDistribute);
79
80} // namespace gpu
81} // namespace transform
82
83namespace gpu {
84void registerTransformDialectExtension(DialectRegistry &registry);
85} // namespace gpu
86} // namespace mlir
87
88#endif // MLIR_DIALECT_GPU_TRANSFORMOPS_GPUTRANSFORMOPS_H
89

source code of mlir/include/mlir/Dialect/GPU/TransformOps/GPUTransformOps.h