Skip to content

Commit

Permalink
[mlir][emitc] Remove copy from scf.for lowering (#94898)
Browse files Browse the repository at this point in the history 
Remove the copy into fresh variables done when lowering scf.for into
emitc.for and use the variables carrying the init and iter values as
the loop's results.
  • Loading branch information
@aniragil
aniragil authored Jun 10, 2024
1 parent fe0dee4 commit 8b7e836
Show file tree
Hide file tree
Showing 3 changed files with 13 additions and 44 deletions.
11 changes: 3 additions & 8 deletions mlir/lib/Conversion/SCFToEmitC/SCFToEmitC.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,10 +102,8 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,
// assigned to by emitc::assign ops within the loop body.
SmallVector<Value> resultVariables =
createVariablesForResults(forOp, rewriter);
SmallVector<Value> iterArgsVariables =
createVariablesForResults(forOp, rewriter);

assignValues(forOp.getInits(), iterArgsVariables, rewriter, loc);
assignValues(forOp.getInits(), resultVariables, rewriter, loc);

emitc::ForOp loweredFor = rewriter.create<emitc::ForOp>(
loc, forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep());
Expand All @@ -117,15 +115,12 @@ LogicalResult ForLowering::matchAndRewrite(ForOp forOp,

SmallVector<Value> replacingValues;
replacingValues.push_back(loweredFor.getInductionVar());
replacingValues.append(iterArgsVariables.begin(), iterArgsVariables.end());
replacingValues.append(resultVariables.begin(), resultVariables.end());

rewriter.mergeBlocks(forOp.getBody(), loweredBody, replacingValues);
lowerYield(iterArgsVariables, rewriter,
lowerYield(resultVariables, rewriter,
cast<scf::YieldOp>(loweredBody->getTerminator()));

// Copy iterArgs into results after the for loop.
assignValues(iterArgsVariables, resultVariables, rewriter, loc);

rewriter.replaceOp(forOp, resultVariables);
return success();
}
Expand Down
28 changes: 10 additions & 18 deletions mlir/test/Conversion/SCFToEmitC/for.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,17 +49,13 @@ func.func @for_yield(%arg0 : index, %arg1 : index, %arg2 : index) -> (f32, f32)
// CHECK-NEXT: %[[VAL_4:.*]] = arith.constant 1.000000e+00 : f32
// CHECK-NEXT: %[[VAL_5:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: %[[VAL_6:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: %[[VAL_7:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: %[[VAL_8:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: emitc.assign %[[VAL_3]] : f32 to %[[VAL_7]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_4]] : f32 to %[[VAL_8]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_3]] : f32 to %[[VAL_5]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_4]] : f32 to %[[VAL_6]] : f32
// CHECK-NEXT: emitc.for %[[VAL_9:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[VAL_2]] {
// CHECK-NEXT: %[[VAL_10:.*]] = arith.addf %[[VAL_7]], %[[VAL_8]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_7]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_8]] : f32
// CHECK-NEXT: %[[VAL_10:.*]] = arith.addf %[[VAL_5]], %[[VAL_6]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_5]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_6]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: emitc.assign %[[VAL_7]] : f32 to %[[VAL_5]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_8]] : f32 to %[[VAL_6]] : f32
// CHECK-NEXT: return %[[VAL_5]], %[[VAL_6]] : f32, f32
// CHECK-NEXT: }

Expand All @@ -78,19 +74,15 @@ func.func @nested_for_yield(%arg0 : index, %arg1 : index, %arg2 : index) -> f32
// CHECK-SAME: %[[VAL_0:.*]]: index, %[[VAL_1:.*]]: index, %[[VAL_2:.*]]: index) -> f32 {
// CHECK-NEXT: %[[VAL_3:.*]] = arith.constant 1.000000e+00 : f32
// CHECK-NEXT: %[[VAL_4:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: %[[VAL_5:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: emitc.assign %[[VAL_3]] : f32 to %[[VAL_5]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_3]] : f32 to %[[VAL_4]] : f32
// CHECK-NEXT: emitc.for %[[VAL_6:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[VAL_2]] {
// CHECK-NEXT: %[[VAL_7:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: %[[VAL_8:.*]] = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
// CHECK-NEXT: emitc.assign %[[VAL_5]] : f32 to %[[VAL_8]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_4]] : f32 to %[[VAL_7]] : f32
// CHECK-NEXT: emitc.for %[[VAL_9:.*]] = %[[VAL_0]] to %[[VAL_1]] step %[[VAL_2]] {
// CHECK-NEXT: %[[VAL_10:.*]] = arith.addf %[[VAL_8]], %[[VAL_8]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_8]] : f32
// CHECK-NEXT: %[[VAL_10:.*]] = arith.addf %[[VAL_7]], %[[VAL_7]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_10]] : f32 to %[[VAL_7]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: emitc.assign %[[VAL_8]] : f32 to %[[VAL_7]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_7]] : f32 to %[[VAL_5]] : f32
// CHECK-NEXT: emitc.assign %[[VAL_7]] : f32 to %[[VAL_4]] : f32
// CHECK-NEXT: }
// CHECK-NEXT: emitc.assign %[[VAL_5]] : f32 to %[[VAL_4]] : f32
// CHECK-NEXT: return %[[VAL_4]] : f32
// CHECK-NEXT: }
18 changes: 0 additions & 18 deletions mlir/test/Target/Cpp/for.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,6 @@ func.func @test_for_yield() {
%s0 = "emitc.constant"() <{value = 0 : i32}> : () -> i32
%p0 = "emitc.constant"() <{value = 1.0 : f32}> : () -> f32

%0 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> i32
%1 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
%2 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> i32
%3 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
emitc.assign %s0 : i32 to %2 : i32
Expand All @@ -53,8 +51,6 @@ func.func @test_for_yield() {
emitc.assign %pn : f32 to %3 : f32
emitc.yield
}
emitc.assign %2 : i32 to %0 : i32
emitc.assign %3 : f32 to %1 : f32

return
}
Expand All @@ -64,8 +60,6 @@ func.func @test_for_yield() {
// CPP-DEFAULT-NEXT: size_t [[STEP:[^ ]*]] = 1;
// CPP-DEFAULT-NEXT: int32_t [[S0:[^ ]*]] = 0;
// CPP-DEFAULT-NEXT: float [[P0:[^ ]*]] = 1.000000000e+00f;
// CPP-DEFAULT-NEXT: int32_t [[SE:[^ ]*]];
// CPP-DEFAULT-NEXT: float [[PE:[^ ]*]];
// CPP-DEFAULT-NEXT: int32_t [[SI:[^ ]*]];
// CPP-DEFAULT-NEXT: float [[PI:[^ ]*]];
// CPP-DEFAULT-NEXT: [[SI:[^ ]*]] = [[S0]];
Expand All @@ -76,8 +70,6 @@ func.func @test_for_yield() {
// CPP-DEFAULT-NEXT: [[SI]] = [[SN]];
// CPP-DEFAULT-NEXT: [[PI]] = [[PN]];
// CPP-DEFAULT-NEXT: }
// CPP-DEFAULT-NEXT: [[SE]] = [[SI]];
// CPP-DEFAULT-NEXT: [[PE]] = [[PI]];
// CPP-DEFAULT-NEXT: return;

// CPP-DECLTOP: void test_for_yield() {
Expand All @@ -86,8 +78,6 @@ func.func @test_for_yield() {
// CPP-DECLTOP-NEXT: size_t [[STEP:[^ ]*]];
// CPP-DECLTOP-NEXT: int32_t [[S0:[^ ]*]];
// CPP-DECLTOP-NEXT: float [[P0:[^ ]*]];
// CPP-DECLTOP-NEXT: int32_t [[SE:[^ ]*]];
// CPP-DECLTOP-NEXT: float [[PE:[^ ]*]];
// CPP-DECLTOP-NEXT: int32_t [[SI:[^ ]*]];
// CPP-DECLTOP-NEXT: float [[PI:[^ ]*]];
// CPP-DECLTOP-NEXT: int32_t [[SN:[^ ]*]];
Expand All @@ -99,8 +89,6 @@ func.func @test_for_yield() {
// CPP-DECLTOP-NEXT: [[P0]] = 1.000000000e+00f;
// CPP-DECLTOP-NEXT: ;
// CPP-DECLTOP-NEXT: ;
// CPP-DECLTOP-NEXT: ;
// CPP-DECLTOP-NEXT: ;
// CPP-DECLTOP-NEXT: [[SI:[^ ]*]] = [[S0]];
// CPP-DECLTOP-NEXT: [[PI:[^ ]*]] = [[P0]];
// CPP-DECLTOP-NEXT: for (size_t [[ITER:[^ ]*]] = [[START]]; [[ITER]] < [[STOP]]; [[ITER]] += [[STEP]]) {
Expand All @@ -109,8 +97,6 @@ func.func @test_for_yield() {
// CPP-DECLTOP-NEXT: [[SI]] = [[SN]];
// CPP-DECLTOP-NEXT: [[PI]] = [[PN]];
// CPP-DECLTOP-NEXT: }
// CPP-DECLTOP-NEXT: [[SE]] = [[SI]];
// CPP-DECLTOP-NEXT: [[PE]] = [[PI]];
// CPP-DECLTOP-NEXT: return;

func.func @test_for_yield_2() {
Expand All @@ -121,8 +107,6 @@ func.func @test_for_yield_2() {
%s0 = emitc.literal "0" : i32
%p0 = emitc.literal "M_PI" : f32

%0 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> i32
%1 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
%2 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> i32
%3 = "emitc.variable"() <{value = #emitc.opaque<"">}> : () -> f32
emitc.assign %s0 : i32 to %2 : i32
Expand All @@ -134,8 +118,6 @@ func.func @test_for_yield_2() {
emitc.assign %pn : f32 to %3 : f32
emitc.yield
}
emitc.assign %2 : i32 to %0 : i32
emitc.assign %3 : f32 to %1 : f32

return
}
Expand Down

0 comments on commit 8b7e836

Please sign in to comment.