From 8c83c3dba275c7f70c2df2fda875cab5fedf5a98 Mon Sep 17 00:00:00 2001 From: Jungwook Park Date: Wed, 30 Jul 2025 14:52:02 +0000 Subject: [PATCH 1/4] [mlir][scf] Add `no_inline` attribute to `scf.execute_region` More control over the IR. Enabling users to explicitly specify which regions should be preserved, uncovers additional opportunities to utilize `scf.execute_region`. --- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 4 ++++ mlir/lib/Dialect/SCF/IR/SCF.cpp | 2 +- mlir/test/Dialect/SCF/canonicalize.mlir | 22 ++++++++++++++++++++++ 3 files changed, 27 insertions(+), 1 deletion(-) diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td index 2d15544e871b3..e6b83fca771a9 100644 --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -119,6 +119,10 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [ ``` }]; + let arguments = (ins + UnitAttr:$no_inline + ); + let results = (outs Variadic); let regions = (region AnyRegion:$region); diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index 759e58b617578..6b7f11a3469d6 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -184,7 +184,7 @@ struct SingleBlockExecuteInliner : public OpRewritePattern { LogicalResult matchAndRewrite(ExecuteRegionOp op, PatternRewriter &rewriter) const override { - if (!op.getRegion().hasOneBlock()) + if (!op.getRegion().hasOneBlock() || op.getNoInline()) return failure(); replaceOpWithRegion(rewriter, op, op.getRegion()); return success(); diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir index 12d30e17f4a8f..ae02144f06f32 100644 --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -1461,6 +1461,28 @@ func.func @execute_region_elim() { // ----- +// CHECK-LABEL: func @execute_region_elim_noinline +func.func @execute_region_elim_noinline() { + affine.for %i = 0 to 100 { + "test.foo"() : () -> () + %v = scf.execute_region -> i64 { + %x = "test.val"() : () -> i64 + scf.yield %x : i64 + } {no_inline} + "test.bar"(%v) : (i64) -> () + } + return +} + +// CHECK-NEXT: affine.for %arg0 = 0 to 100 { +// CHECK-NEXT: "test.foo"() : () -> () +// CHECK-NEXT: scf.execute_region +// CHECK-NEXT: %[[VAL:.*]] = "test.val"() : () -> i64 +// CHECK-NEXT: scf.yield %[[VAL]] : i64 +// CHECK-NEXT: } + +// ----- + // CHECK-LABEL: func @func_execute_region_elim func.func @func_execute_region_elim() { "test.foo"() : () -> () From 21c92a26714c6882ea28cfd2b1add9193dff4e06 Mon Sep 17 00:00:00 2001 From: Jungwook Park Date: Wed, 30 Jul 2025 19:52:36 +0000 Subject: [PATCH 2/4] Let attribute as a keyword --- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 11 +++++++++++ mlir/lib/Dialect/SCF/IR/SCF.cpp | 6 +++++- mlir/test/Dialect/SCF/canonicalize.mlir | 20 ++++++++++---------- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td index e6b83fca771a9..8cf6895bc3dfd 100644 --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -87,6 +87,9 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [ be accessed inside the op. The op's region can have multiple blocks and the blocks can have multiple distinct terminators. Values returned from this op's region define the op's results. + Canonicalizer inlines an ExecuteRegionOp into its parent if it only contains + one block or its parent can contain multiple blocks, 'no_inline' attribute + can be set to prevent an ExecuteRegionOp from being inlined. Example: @@ -98,6 +101,14 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [ } } + // the same as above but with no_inline attribute + scf.for %i = 0 to 128 step %c1 { + %y = scf.execute_region -> i32 no_inline { + %x = load %A[%i] : memref<128xi32> + scf.yield %x : i32 + } + } + affine.for %i = 0 to 100 { "foo"() : () -> () %v = scf.execute_region -> i64 { diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index 6b7f11a3469d6..59d7ee473a2cc 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -137,6 +137,9 @@ ParseResult ExecuteRegionOp::parse(OpAsmParser &parser, if (parser.parseOptionalArrowTypeList(result.types)) return failure(); + if (succeeded(parser.parseOptionalKeyword("no_inline"))) + result.addAttribute("no_inline", parser.getBuilder().getUnitAttr()); + // Introduce the body region and parse it. Region *body = result.addRegion(); if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}) || @@ -148,8 +151,9 @@ ParseResult ExecuteRegionOp::parse(OpAsmParser &parser, void ExecuteRegionOp::print(OpAsmPrinter &p) { p.printOptionalArrowTypeList(getResultTypes()); - p << ' '; + if(getNoInline()) + p << "no_inline "; p.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true); diff --git a/mlir/test/Dialect/SCF/canonicalize.mlir b/mlir/test/Dialect/SCF/canonicalize.mlir index ae02144f06f32..308cf150aa98e 100644 --- a/mlir/test/Dialect/SCF/canonicalize.mlir +++ b/mlir/test/Dialect/SCF/canonicalize.mlir @@ -1440,8 +1440,8 @@ func.func @propagate_into_execute_region() { // ----- -// CHECK-LABEL: func @execute_region_elim -func.func @execute_region_elim() { +// CHECK-LABEL: func @execute_region_inline +func.func @execute_region_inline() { affine.for %i = 0 to 100 { "test.foo"() : () -> () %v = scf.execute_region -> i64 { @@ -1461,14 +1461,14 @@ func.func @execute_region_elim() { // ----- -// CHECK-LABEL: func @execute_region_elim_noinline -func.func @execute_region_elim_noinline() { +// CHECK-LABEL: func @execute_region_no_inline +func.func @execute_region_no_inline() { affine.for %i = 0 to 100 { "test.foo"() : () -> () - %v = scf.execute_region -> i64 { + %v = scf.execute_region -> i64 no_inline { %x = "test.val"() : () -> i64 scf.yield %x : i64 - } {no_inline} + } "test.bar"(%v) : (i64) -> () } return @@ -1483,8 +1483,8 @@ func.func @execute_region_elim_noinline() { // ----- -// CHECK-LABEL: func @func_execute_region_elim -func.func @func_execute_region_elim() { +// CHECK-LABEL: func @func_execute_region_inline +func.func @func_execute_region_inline() { "test.foo"() : () -> () %v = scf.execute_region -> i64 { %c = "test.cmp"() : () -> i1 @@ -1518,8 +1518,8 @@ func.func @func_execute_region_elim() { // ----- -// CHECK-LABEL: func @func_execute_region_elim_multi_yield -func.func @func_execute_region_elim_multi_yield() { +// CHECK-LABEL: func @func_execute_region_inline_multi_yield +func.func @func_execute_region_inline_multi_yield() { "test.foo"() : () -> () %v = scf.execute_region -> i64 { %c = "test.cmp"() : () -> i1 From 3564d75f93a509edae102a02a72b17a6bafb7d9f Mon Sep 17 00:00:00 2001 From: Jungwook Park Date: Wed, 30 Jul 2025 19:58:51 +0000 Subject: [PATCH 3/4] format --- mlir/lib/Dialect/SCF/IR/SCF.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mlir/lib/Dialect/SCF/IR/SCF.cpp b/mlir/lib/Dialect/SCF/IR/SCF.cpp index 59d7ee473a2cc..0262a1b8a3893 100644 --- a/mlir/lib/Dialect/SCF/IR/SCF.cpp +++ b/mlir/lib/Dialect/SCF/IR/SCF.cpp @@ -152,7 +152,7 @@ ParseResult ExecuteRegionOp::parse(OpAsmParser &parser, void ExecuteRegionOp::print(OpAsmPrinter &p) { p.printOptionalArrowTypeList(getResultTypes()); p << ' '; - if(getNoInline()) + if (getNoInline()) p << "no_inline "; p.printRegion(getRegion(), /*printEntryBlockArgs=*/false, From 03a3c0bea3a56e068ac247b311ccd4e81fb388c5 Mon Sep 17 00:00:00 2001 From: Jungwook Park Date: Thu, 31 Jul 2025 16:46:50 +0000 Subject: [PATCH 4/4] Revised dosument from the review. --- mlir/include/mlir/Dialect/SCF/IR/SCFOps.td | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td index 8cf6895bc3dfd..0c1c15b85f4c9 100644 --- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td +++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td @@ -87,9 +87,9 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [ be accessed inside the op. The op's region can have multiple blocks and the blocks can have multiple distinct terminators. Values returned from this op's region define the op's results. - Canonicalizer inlines an ExecuteRegionOp into its parent if it only contains - one block or its parent can contain multiple blocks, 'no_inline' attribute - can be set to prevent an ExecuteRegionOp from being inlined. + The optional 'no_inline' flag can be set to request the ExecuteRegionOp to be + preserved as much as possible and not being inlined in the parent block until + an explicit lowering step. Example: