[mlir]introduce UnrollScopeInterface and apply it to funcOp and gpu.launch Op.

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;

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 = [{

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 {

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<InferIntRangeInterface, ["inferResultRanges"]>,
-      RecursiveMemoryEffects]>,
+      RecursiveMemoryEffects, UnrollScopeInterface]>,
     Arguments<(ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                Index:$gridSizeX, Index:$gridSizeY, Index:$gridSizeZ,
                Index:$blockSizeX, Index:$blockSizeY, Index:$blockSizeZ,

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)

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_

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

mlir/lib/Dialect/Affine/Utils/CMakeLists.txt

@@ -16,4 +16,5 @@ add_mlir_dialect_library(MLIRAffineUtils
   MLIRMemRefDialect
   MLIRTransformUtils
   MLIRViewLikeInterface
+  MLIRUnrollScopeInterface
   )

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<FunctionOpInterface>();
+      auto unrollScope = forOp->getParentOfType<UnrollScopeInterface>();
       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<arith::ConstantIndexOp>(

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

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"

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 {