From 9233d6369a19dfd42370e87a220de92f621926a6 Mon Sep 17 00:00:00 2001 From: linuxlonelyeagle <2020382038@qq.com> Date: Wed, 22 Jan 2025 15:44:09 +0800 Subject: [PATCH] introduce UnrollScopeInterface and apply it to funcOp and gpu.launch Op. --- mlir/include/mlir/Dialect/Func/IR/FuncOps.h | 1 + mlir/include/mlir/Dialect/Func/IR/FuncOps.td | 5 ++- mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h | 1 + mlir/include/mlir/Dialect/GPU/IR/GPUOps.td | 3 +- mlir/include/mlir/Interfaces/CMakeLists.txt | 1 + .../mlir/Interfaces/UnrollScopeInterface.h | 21 ++++++++++ .../mlir/Interfaces/UnrollScopeInterface.td | 36 +++++++++++++++++ mlir/lib/Dialect/Affine/Utils/CMakeLists.txt | 1 + mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 7 ++-- mlir/lib/Interfaces/CMakeLists.txt | 2 + mlir/lib/Interfaces/UnrollScopeInterface.cpp | 18 +++++++++ mlir/test/Dialect/Affine/unroll.mlir | 40 +++++++++++++++++++ 12 files changed, 130 insertions(+), 6 deletions(-) create mode 100644 mlir/include/mlir/Interfaces/UnrollScopeInterface.h create mode 100644 mlir/include/mlir/Interfaces/UnrollScopeInterface.td create mode 100644 mlir/lib/Interfaces/UnrollScopeInterface.cpp diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.h b/mlir/include/mlir/Dialect/Func/IR/FuncOps.h index 5e10a9f50b774..3f5566a28546d 100644 --- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.h +++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.h @@ -20,6 +20,7 @@ #include "mlir/Interfaces/FunctionInterfaces.h" #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Interfaces/UnrollScopeInterface.h" namespace mlir { class PatternRewriter; diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td index 4da0efcb13ddf..5c9f8c6a59f8f 100644 --- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td +++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td @@ -17,6 +17,7 @@ include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/FunctionInterfaces.td" include "mlir/Interfaces/InferTypeOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/UnrollScopeInterface.td" def Func_Dialect : Dialect { let name = "func"; @@ -225,8 +226,8 @@ def ConstantOp : Func_Op<"constant", //===----------------------------------------------------------------------===// def FuncOp : Func_Op<"func", [ - AffineScope, AutomaticAllocationScope, - FunctionOpInterface, IsolatedFromAbove, OpAsmOpInterface + AffineScope, AutomaticAllocationScope, FunctionOpInterface, + IsolatedFromAbove, OpAsmOpInterface, UnrollScopeInterface ]> { let summary = "An operation with a name containing a single `SSACFG` region"; let description = [{ diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h index 7b53594a1c8e2..0cf2d0c77383f 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h @@ -29,6 +29,7 @@ #include "mlir/Interfaces/InferIntRangeInterface.h" #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Interfaces/UnrollScopeInterface.h" #include "llvm/ADT/STLExtras.h" namespace mlir { diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td index 3adfd5f4f2c43..8279bb9985ea3 100644 --- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td +++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td @@ -30,6 +30,7 @@ include "mlir/Interfaces/FunctionInterfaces.td" include "mlir/Interfaces/InferIntRangeInterface.td" include "mlir/Interfaces/InferTypeOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/UnrollScopeInterface.td" //===----------------------------------------------------------------------===// // GPU Dialect operations. @@ -796,7 +797,7 @@ def GPU_LaunchFuncOp :GPU_Op<"launch_func", [ def GPU_LaunchOp : GPU_Op<"launch", [ AutomaticAllocationScope, AttrSizedOperandSegments, GPU_AsyncOpInterface, DeclareOpInterfaceMethods, - RecursiveMemoryEffects]>, + RecursiveMemoryEffects, UnrollScopeInterface]>, Arguments<(ins Variadic:$asyncDependencies, Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ, Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ, diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt index d81298bb4daf0..cd6cc084dd280 100644 --- a/mlir/include/mlir/Interfaces/CMakeLists.txt +++ b/mlir/include/mlir/Interfaces/CMakeLists.txt @@ -17,6 +17,7 @@ add_mlir_interface(TilingInterface) add_mlir_interface(ValueBoundsOpInterface) add_mlir_interface(VectorInterfaces) add_mlir_interface(ViewLikeInterface) +add_mlir_interface(UnrollScopeInterface) set(LLVM_TARGET_DEFINITIONS MemorySlotInterfaces.td) mlir_tablegen(MemorySlotOpInterfaces.h.inc -gen-op-interface-decls) diff --git a/mlir/include/mlir/Interfaces/UnrollScopeInterface.h b/mlir/include/mlir/Interfaces/UnrollScopeInterface.h new file mode 100644 index 0000000000000..f7d71b6f9be65 --- /dev/null +++ b/mlir/include/mlir/Interfaces/UnrollScopeInterface.h @@ -0,0 +1,21 @@ +//===- UnrollScopeInterface.h - unroll region interface -------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// +// +// This file implements the operation interface for unroll region +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_UNROLLSCOPEINTERFACE_H_ +#define MLIR_INTERFACES_UNROLLSCOPEINTERFACE_H_ + +#include "mlir/IR/OpDefinition.h" + +/// Include the generated interface declarations. +#include "mlir/Interfaces/UnrollScopeInterface.h.inc" + +#endif // MLIR_INTERFACES_UNROLLSCOPEINTERFACE_H_ diff --git a/mlir/include/mlir/Interfaces/UnrollScopeInterface.td b/mlir/include/mlir/Interfaces/UnrollScopeInterface.td new file mode 100644 index 0000000000000..5ad5e5b44cfe1 --- /dev/null +++ b/mlir/include/mlir/Interfaces/UnrollScopeInterface.td @@ -0,0 +1,36 @@ +//===- UnrollScopeInterface.td - unroll scope interface ----*- 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 +// +//===----------------------------------------------------------------------===// +// +// Defines the interface for unroll region. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_UNROLLSCOPEINTERFACE +#define MLIR_INTERFACES_UNROLLSCOPEINTERFACE + +include "mlir/IR/OpBase.td" + +def UnrollScopeInterface : OpInterface<"UnrollScopeInterface"> { + let description = [{ + This interface controls the scope of the loop unroll.It ensures + that SSA values generated outside the loop when unrolling are + in the nearest `UnrollScopeInterface` region. + }]; + let cppNamespace = "::mlir"; + let methods = [ + InterfaceMethod<[{ + return the `UnrollScopeInterface` region. + }], + "::mlir::Region&", "getUnrollBody", (ins), + /*methodBody=*/[{}], /*defaultImplementation=*/[{ + return $_op->getRegion(0); + }]>, + ]; +} + +#endif // MLIR_INTERFACES_UNROLLSCOPEINTERFACE diff --git a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt index ef6e0dbf45d3a..d2993a424060f 100644 --- a/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt +++ b/mlir/lib/Dialect/Affine/Utils/CMakeLists.txt @@ -16,4 +16,5 @@ add_mlir_dialect_library(MLIRAffineUtils MLIRMemRefDialect MLIRTransformUtils MLIRViewLikeInterface + MLIRUnrollScopeInterface ) diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp index 4e02559a08949..c697d3b0127a8 100644 --- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp @@ -21,6 +21,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/IntegerSet.h" +#include "mlir/Interfaces/UnrollScopeInterface.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "llvm/ADT/MapVector.h" #include "llvm/ADT/SmallPtrSet.h" @@ -129,10 +130,10 @@ LogicalResult mlir::affine::promoteIfSingleIteration(AffineForOp forOp) { auto *parentBlock = forOp->getBlock(); if (!iv.use_empty()) { if (forOp.hasConstantLowerBound()) { - auto func = forOp->getParentOfType(); + auto unrollScope = forOp->getParentOfType(); OpBuilder builder(forOp->getContext()); - if (func) - builder.setInsertionPointToStart(&func.getFunctionBody().front()); + if (unrollScope) + builder.setInsertionPointToStart(&unrollScope.getUnrollBody().front()); else builder.setInsertionPoint(forOp); auto constOp = builder.create( diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt index d3b7bf65ad3e7..bd0b79aaecab0 100644 --- a/mlir/lib/Interfaces/CMakeLists.txt +++ b/mlir/lib/Interfaces/CMakeLists.txt @@ -21,6 +21,7 @@ set(LLVM_OPTIONAL_SOURCES ValueBoundsOpInterface.cpp VectorInterfaces.cpp ViewLikeInterface.cpp + UnrollScopeInterface.cpp ) function(add_mlir_interface_library name) @@ -46,6 +47,7 @@ add_mlir_interface_library(CopyOpInterface) add_mlir_interface_library(DataLayoutInterfaces) add_mlir_interface_library(DerivedAttributeOpInterface) add_mlir_interface_library(DestinationStyleOpInterface) +add_mlir_interface_library(UnrollScopeInterface) add_mlir_library(MLIRFunctionInterfaces FunctionInterfaces.cpp diff --git a/mlir/lib/Interfaces/UnrollScopeInterface.cpp b/mlir/lib/Interfaces/UnrollScopeInterface.cpp new file mode 100644 index 0000000000000..b500f5ad01f2a --- /dev/null +++ b/mlir/lib/Interfaces/UnrollScopeInterface.cpp @@ -0,0 +1,18 @@ +//===- UnrollScopeInterface.cpp - unroll scope interface in MLIR ----------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include "mlir/Interfaces/UnrollScopeInterface.h" + +using namespace mlir; + +//===----------------------------------------------------------------------===// +// UnrollScopeInterface Interface +//===----------------------------------------------------------------------===// + +/// Include the definitions of the unroll scope interface. +#include "mlir/Interfaces/UnrollScopeInterface.cpp.inc" diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir index e398c3fe2011d..4492fdf10515f 100644 --- a/mlir/test/Dialect/Affine/unroll.mlir +++ b/mlir/test/Dialect/Affine/unroll.mlir @@ -240,6 +240,46 @@ func.func @loop_nest_unroll_full() { return } // UNROLL-FULL } + +// UNROLL-FULL-LABEL: func @gpu_launch_unroll + +func.func @gpu_launch_unroll() { + %c1 = arith.constant 1 : index + gpu.launch blocks(%arg0, %arg1, %arg2) in (%arg6 = %c1, %arg7 = %c1, %arg8 = %c1) threads(%arg3, %arg4, %arg5) in (%arg9 = %c1, %arg10 = %c1, %arg11 = %c1) { + %cst = arith.constant dense<0.000000e+00> : vector<2x4x2x2xf16> + %0 = affine.for %arg12 = 0 to 2 iter_args(%arg13 = %cst) -> (vector<2x4x2x2xf16>) { + %1 = affine.for %arg14 = 0 to 4 iter_args(%arg15 = %arg13) -> (vector<2x4x2x2xf16>) { + %cst_0 = arith.constant dense<0.000000e+00> : vector<2x2xf16> + %2 = vector.insert %cst_0, %arg13 [%arg12, %arg14] : vector<2x2xf16> into vector<2x4x2x2xf16> + affine.yield %2 : vector<2x4x2x2xf16> + } + affine.yield %1 : vector<2x4x2x2xf16> + } + gpu.terminator + } + return +} + +// UNROLL-FULL: %[[VAL_0:.*]] = arith.constant 1 : index +// UNROLL-FULL: gpu.launch blocks(%[[VAL_1:.*]], %[[VAL_2:.*]], %[[VAL_3:.*]]) in (%[[VAL_4:.*]] = %[[VAL_0]], %[[VAL_5:.*]] = %[[VAL_0]], %[[VAL_6:.*]] = %[[VAL_0]]) threads(%[[VAL_7:.*]], %[[VAL_8:.*]], %[[VAL_9:.*]]) in (%[[VAL_10:.*]] = %[[VAL_0]], %[[VAL_11:.*]] = %[[VAL_0]], %[[VAL_12:.*]] = %[[VAL_0]]) { +// UNROLL-FULL: %[[VAL_13:.*]] = arith.constant 0 : index +// UNROLL-FULL: %[[VAL_14:.*]] = arith.constant dense<0.000000e+00> : vector<2x4x2x2xf16> +// UNROLL-FULL: %[[VAL_15:.*]] = affine.for %[[VAL_16:.*]] = 0 to 2 iter_args(%[[VAL_17:.*]] = %[[VAL_14]]) -> (vector<2x4x2x2xf16>) { +// UNROLL-FULL: %[[VAL_18:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf16> +// UNROLL-FULL: %[[VAL_19:.*]] = vector.insert %[[VAL_18]], %[[VAL_17]] {{\[}}%[[VAL_16]], %[[VAL_13]]] : vector<2x2xf16> into vector<2x4x2x2xf16> +// UNROLL-FULL: %[[VAL_20:.*]] = affine.apply [[$MAP0]](%[[VAL_13]]) +// UNROLL-FULL: %[[VAL_21:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf16> +// UNROLL-FULL: %[[VAL_22:.*]] = vector.insert %[[VAL_21]], %[[VAL_17]] {{\[}}%[[VAL_16]], %[[VAL_20]]] : vector<2x2xf16> into vector<2x4x2x2xf16> +// UNROLL-FULL: %[[VAL_23:.*]] = affine.apply [[$MAP1]](%[[VAL_13]]) +// UNROLL-FULL: %[[VAL_24:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf16> +// UNROLL-FULL: %[[VAL_25:.*]] = vector.insert %[[VAL_24]], %[[VAL_17]] {{\[}}%[[VAL_16]], %[[VAL_23]]] : vector<2x2xf16> into vector<2x4x2x2xf16> +// UNROLL-FULL: %[[VAL_26:.*]] = affine.apply [[$MAP2]](%[[VAL_13]]) +// UNROLL-FULL: %[[VAL_27:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf16> +// UNROLL-FULL: %[[VAL_28:.*]] = vector.insert %[[VAL_27]], %[[VAL_17]] {{\[}}%[[VAL_16]], %[[VAL_26]]] : vector<2x2xf16> into vector<2x4x2x2xf16> +// UNROLL-FULL: affine.yield %[[VAL_28]] : vector<2x4x2x2xf16> +// UNROLL-FULL: } +// UNROLL-FULL: gpu.terminator + // SHORT-LABEL: func @loop_nest_outer_unroll() { func.func @loop_nest_outer_unroll() { // SHORT: affine.for %arg0 = 0 to 4 {