diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h index bc29d04287ac4..ea5034b60d8bd 100644 --- a/mlir/include/mlir/Dialect/Affine/Passes.h +++ b/mlir/include/mlir/Dialect/Affine/Passes.h @@ -14,6 +14,7 @@ #ifndef MLIR_DIALECT_AFFINE_PASSES_H #define MLIR_DIALECT_AFFINE_PASSES_H +#include "mlir/Interfaces/FunctionInterfaces.h" #include "mlir/Pass/Pass.h" #include @@ -93,7 +94,7 @@ std::unique_ptr> createLoopTilingPass(); /// factors supplied through other means. If -1 is passed as the unrollFactor /// and no callback is provided, anything passed from the command-line (if at /// all) or the default unroll factor is used (LoopUnroll:kDefaultUnrollFactor). -std::unique_ptr> createLoopUnrollPass( +std::unique_ptr> createLoopUnrollPass( int unrollFactor = -1, bool unrollUpToFactor = false, bool unrollFull = false, const std::function &getUnrollFactor = nullptr); diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td index d7c7897c65730..5325d3b0a1d69 100644 --- a/mlir/include/mlir/Dialect/Affine/Passes.td +++ b/mlir/include/mlir/Dialect/Affine/Passes.td @@ -199,7 +199,7 @@ def AffineLoopTiling : Pass<"affine-loop-tile", "func::FuncOp"> { ]; } -def AffineLoopUnroll : Pass<"affine-loop-unroll", "func::FuncOp"> { +def AffineLoopUnroll : InterfacePass<"affine-loop-unroll", "FunctionOpInterface"> { let summary = "Unroll affine loops"; let constructor = "mlir::affine::createLoopUnrollPass()"; let options = [ diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp index 57df7ada91654..7ff77968c61ad 100644 --- a/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/LoopUnroll.cpp @@ -82,7 +82,7 @@ static bool isInnermostAffineForOp(AffineForOp op) { } /// Gathers loops that have no affine.for's nested within. -static void gatherInnermostLoops(func::FuncOp f, +static void gatherInnermostLoops(FunctionOpInterface f, SmallVectorImpl &loops) { f.walk([&](AffineForOp forOp) { if (isInnermostAffineForOp(forOp)) @@ -91,7 +91,7 @@ static void gatherInnermostLoops(func::FuncOp f, } void LoopUnroll::runOnOperation() { - func::FuncOp func = getOperation(); + FunctionOpInterface func = getOperation(); if (func.isExternal()) return; @@ -100,8 +100,8 @@ void LoopUnroll::runOnOperation() { SmallVector loops; // Gathers all loops with trip count <= minTripCount. Do a post order walk - // so that loops are gathered from innermost to outermost (or else unrolling - // an outer one may delete gathered inner ones). + // so that loops are gathered from innermost to outermost (or else + // unrolling an outer one may delete gathered inner ones). getOperation().walk([&](AffineForOp forOp) { std::optional tripCount = getConstantTripCount(forOp); if (tripCount && *tripCount <= unrollFullThreshold) @@ -145,7 +145,8 @@ LogicalResult LoopUnroll::runOnAffineForOp(AffineForOp forOp) { cleanUpUnroll); } -std::unique_ptr> mlir::affine::createLoopUnrollPass( +std::unique_ptr> +mlir::affine::createLoopUnrollPass( int unrollFactor, bool unrollUpToFactor, bool unrollFull, const std::function &getUnrollFactor) { return std::make_unique( diff --git a/mlir/test/Dialect/Affine/unroll.mlir b/mlir/test/Dialect/Affine/unroll.mlir index e398c3fe2011d..574e9f41494af 100644 --- a/mlir/test/Dialect/Affine/unroll.mlir +++ b/mlir/test/Dialect/Affine/unroll.mlir @@ -1,8 +1,9 @@ -// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-full" | FileCheck %s --check-prefix UNROLL-FULL -// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-full unroll-full-threshold=2" | FileCheck %s --check-prefix SHORT -// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=4" | FileCheck %s --check-prefix UNROLL-BY-4 -// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=1" | FileCheck %s --check-prefix UNROLL-BY-1 -// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-unroll="unroll-factor=5 cleanup-unroll=true" | FileCheck %s --check-prefix UNROLL-CLEANUP-LOOP +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true}))" | FileCheck %s --check-prefix UNROLL-FULL +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-full=true unroll-full-threshold=2}))" | FileCheck %s --check-prefix SHORT +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=4}))" | FileCheck %s --check-prefix UNROLL-BY-4 +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=1}))" | FileCheck %s --check-prefix UNROLL-BY-1 +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix UNROLL-CLEANUP-LOOP +// RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline="builtin.module(gpu.module(gpu.func(affine-loop-unroll{unroll-full=true})))" | FileCheck %s --check-prefix GPU-UNROLL-FULL // UNROLL-FULL-DAG: [[$MAP0:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)> // UNROLL-FULL-DAG: [[$MAP1:#map[0-9]*]] = affine_map<(d0) -> (d0 + 2)> @@ -240,6 +241,23 @@ func.func @loop_nest_unroll_full() { return } // UNROLL-FULL } +gpu.module @unroll_full { + // GPU-UNROLL-FULL-LABEL: func @gpu_loop_nest_simplest() { + gpu.func @gpu_loop_nest_simplest() { + // GPU-UNROLL-FULL: affine.for %arg0 = 0 to 100 step 2 { + affine.for %i = 0 to 100 step 2 { + // GPU-UNROLL-FULL: %c1_i32 = arith.constant 1 : i32 + // GPU-UNROLL-FULL-NEXT: %c1_i32_0 = arith.constant 1 : i32 + // GPU-UNROLL-FULL-NEXT: %c1_i32_1 = arith.constant 1 : i32 + // GPU-UNROLL-FULL-NEXT: %c1_i32_2 = arith.constant 1 : i32 + affine.for %j = 0 to 4 { + %x = arith.constant 1 : i32 + } + } // GPU-UNROLL-FULL: } + gpu.return // GPU-UNROLL-FULL: return + } +} + // SHORT-LABEL: func @loop_nest_outer_unroll() { func.func @loop_nest_outer_unroll() { // SHORT: affine.for %arg0 = 0 to 4 { diff --git a/mlir/test/Dialect/SCF/loop-unroll.mlir b/mlir/test/Dialect/SCF/loop-unroll.mlir index 0368505a1b70d..4c72d9e99d049 100644 --- a/mlir/test/Dialect/SCF/loop-unroll.mlir +++ b/mlir/test/Dialect/SCF/loop-unroll.mlir @@ -3,9 +3,9 @@ // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=0' | FileCheck %s --check-prefix UNROLL-OUTER-BY-2 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 loop-depth=1' | FileCheck %s --check-prefix UNROLL-INNER-BY-2 // RUN: mlir-opt %s -test-loop-unrolling='unroll-factor=2 annotate=true' | FileCheck %s --check-prefix UNROLL-BY-2-ANNOTATE -// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=6 unroll-up-to-factor=true' | FileCheck %s --check-prefix UNROLL-UP-TO -// RUN: mlir-opt %s --affine-loop-unroll='unroll-factor=5 cleanup-unroll=true' | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5 -// RUN: mlir-opt %s --affine-loop-unroll --split-input-file | FileCheck %s +// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=6 unroll-up-to-factor=true}))" | FileCheck %s --check-prefix UNROLL-UP-TO +// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll{unroll-factor=5 cleanup-unroll=true}))" | FileCheck %s --check-prefix CLEANUP-UNROLL-BY-5 +// RUN: mlir-opt %s -pass-pipeline="builtin.module(func.func(affine-loop-unroll))" --split-input-file | FileCheck %s func.func @dynamic_loop_unroll(%arg0 : index, %arg1 : index, %arg2 : index, %arg3: memref) {