Skip to content

Commit ffc9add

Browse files
vimarsh6739vimarsh6739
committed
vibe activity analysis printer
1 parent 8758bd6 commit ffc9add

3 files changed

Lines changed: 87 additions & 24 deletions

File tree

‎enzyme/Enzyme/MLIR/Passes/EnzymeAcitivityOpt.cpp‎

Lines changed: 68 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
#include "mlir/Support/LLVM.h"
2323
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
2424

25+
#include "llvm/Support/raw_ostream.h"
26+
2527
#include <deque>
2628

2729
#define DEBUG_TYPE "enzyme-activity-opt"
@@ -40,9 +42,15 @@ namespace enzyme {
4042

4143
namespace {
4244

45+
struct TaggedActivity {
46+
Attribute tag;
47+
bool isConstant;
48+
};
49+
4350
struct ActivityOptimization {
4451
SmallVector<Activity> argumentActivity;
4552
SmallVector<Activity> resultActivity;
53+
SmallVector<TaggedActivity> taggedActivities;
4654
};
4755

4856
struct ForwardResultSlot {
@@ -140,16 +148,28 @@ static ArrayAttr getActivityArrayAttr(MLIRContext *context,
140148
return ArrayAttr::get(context, attrs);
141149
}
142150

143-
//TODO: to be implemented by initializing an ActivityAnalyzer obj for callee
151+
static void
152+
printTaggedActivityAnalysis(FunctionOpInterface callee,
153+
ArrayRef<TaggedActivity> taggedActivities) {
154+
llvm::errs() << FlatSymbolRefAttr::get(callee) << ":\n";
155+
for (const TaggedActivity &activity : taggedActivities) {
156+
llvm::errs() << " " << activity.tag << ": "
157+
<< (activity.isConstant ? "Constant" : "Active") << "\n";
158+
}
159+
}
160+
161+
// TODO: to be implemented by initializing an ActivityAnalyzer obj for callee
144162
static FailureOr<ActivityOptimization> computeClassicalActivityOptimization(
145163
Operation *diffOp, FunctionOpInterface callee,
146164
ArrayRef<Activity> argumentActivity, ArrayRef<Activity> resultActivity,
147-
SymbolTableCollection &symbolTable) {return failure();}
165+
SymbolTableCollection &symbolTable) {
166+
return failure();
167+
}
148168

149169
static FailureOr<ActivityOptimization> computeDataFlowActivityOptimization(
150170
Operation *diffOp, FunctionOpInterface callee,
151171
ArrayRef<Activity> argumentActivity, ArrayRef<Activity> resultActivity,
152-
SymbolTableCollection &symbolTable) {
172+
SymbolTableCollection &symbolTable, bool collectActivityAnalysis = false) {
153173
Region &body = callee.getFunctionBody();
154174
if (body.empty() || body.front().empty())
155175
return failure();
@@ -201,12 +221,14 @@ static FailureOr<ActivityOptimization> computeDataFlowActivityOptimization(
201221
returnOps.push_back(op);
202222
for (const auto &[operand, activity] :
203223
llvm::zip(op->getOperands(), resultActivity)) {
204-
// we should only seed active/activenn inputs(dup/dupnoneed is handled
205-
// in enzyme::BackwardMemoryActivity)
224+
// Seed scalar return activity here; dense memory activity handles
225+
// pointed-to data for duplicated memory returns.
206226
auto *returnLattice =
207227
solver.getOrCreateState<enzyme::BackwardValueActivity>(operand);
208-
if (activity == Activity::enzyme_active &&
209-
activity == Activity::enzyme_activenoneed) {
228+
if (activity == Activity::enzyme_active ||
229+
activity == Activity::enzyme_activenoneed ||
230+
activity == Activity::enzyme_dup ||
231+
activity == Activity::enzyme_dupnoneed) {
210232
(void)returnLattice->meet(ValueActivity::getActiveVal());
211233
}
212234
}
@@ -305,17 +327,20 @@ static FailureOr<ActivityOptimization> computeDataFlowActivityOptimization(
305327
return !isActiveData(value);
306328
};
307329

330+
// Initialize activity results to be the same as that of the callsite
308331
ActivityOptimization optimized;
309332
optimized.argumentActivity.assign(argumentActivity.begin(),
310333
argumentActivity.end());
311334
optimized.resultActivity.assign(resultActivity.begin(), resultActivity.end());
312335

313336
for (const auto &[index, arg, activity] :
314337
llvm::enumerate(callee.getArguments(), argumentActivity)) {
315-
if (canDemoteArgumentActivity(activity) && isConstantValue(arg))
338+
if (isConstantValue(arg))
316339
optimized.argumentActivity[index] = demoteArgumentActivity(activity);
317340
}
318341

342+
// Conservatively combine all the activities of all ReturnLike ops as a
343+
// summary for ret_activity originating from the callsite autodiff op
319344
SmallVector<bool> allReturnValuesConstant(resultActivity.size(), true);
320345
for (Operation *returnOp : returnOps) {
321346
for (const auto &[index, operand] :
@@ -330,6 +355,21 @@ static FailureOr<ActivityOptimization> computeDataFlowActivityOptimization(
330355
optimized.resultActivity[index] = demoteReturnActivity(activity);
331356
}
332357

358+
if (collectActivityAnalysis) {
359+
for (BlockArgument arg : callee.getArguments()) {
360+
if (Attribute tag = callee.getArgAttr(arg.getArgNumber(), "enzyme.tag"))
361+
optimized.taggedActivities.push_back({tag, isConstantValue(arg)});
362+
}
363+
364+
callee.walk([&](Operation *op) {
365+
Attribute tag = op->getAttr("tag");
366+
if (!tag)
367+
return;
368+
for (OpResult result : op->getResults())
369+
optimized.taggedActivities.push_back({tag, isConstantValue(result)});
370+
});
371+
}
372+
333373
return optimized;
334374
}
335375

@@ -638,8 +678,10 @@ static LogicalResult mapReverseResults(PatternRewriter &rewriter,
638678
class MinimizeForwardActivity : public OpRewritePattern<ForwardDiffOp> {
639679
public:
640680
MinimizeForwardActivity(MLIRContext *context, bool dataflow,
681+
bool printActivityAnalysis,
641682
PatternBenefit benefit = 1)
642-
: OpRewritePattern(context, benefit), dataflow(dataflow) {}
683+
: OpRewritePattern(context, benefit), dataflow(dataflow),
684+
printActivityAnalysis(printActivityAnalysis) {}
643685

644686
LogicalResult matchAndRewrite(ForwardDiffOp op,
645687
PatternRewriter &rewriter) const override {
@@ -657,7 +699,8 @@ class MinimizeForwardActivity : public OpRewritePattern<ForwardDiffOp> {
657699
FailureOr<ActivityOptimization> optimized =
658700
dataflow
659701
? computeDataFlowActivityOptimization(op, callee, argumentActivity,
660-
resultActivity, symbolTable)
702+
resultActivity, symbolTable,
703+
printActivityAnalysis)
661704
: computeClassicalActivityOptimization(op, callee, argumentActivity,
662705
resultActivity, symbolTable);
663706

@@ -668,6 +711,9 @@ class MinimizeForwardActivity : public OpRewritePattern<ForwardDiffOp> {
668711
llvm::equal(resultActivity, optimized->resultActivity))
669712
return failure();
670713

714+
if (printActivityAnalysis && dataflow)
715+
printTaggedActivityAnalysis(callee, optimized->taggedActivities);
716+
671717
FailureOr<SmallVector<Value>> newInputs =
672718
buildForwardInputs(op, argumentActivity, optimized->argumentActivity);
673719
FailureOr<SmallVector<ForwardResultSlot>> oldSlots =
@@ -707,13 +753,16 @@ class MinimizeForwardActivity : public OpRewritePattern<ForwardDiffOp> {
707753

708754
private:
709755
bool dataflow;
756+
bool printActivityAnalysis;
710757
};
711758

712759
class MinimizeReverseActivity : public OpRewritePattern<AutoDiffOp> {
713760
public:
714761
MinimizeReverseActivity(MLIRContext *context, bool dataflow,
762+
bool printActivityAnalysis,
715763
PatternBenefit benefit = 1)
716-
: OpRewritePattern(context, benefit), dataflow(dataflow) {}
764+
: OpRewritePattern(context, benefit), dataflow(dataflow),
765+
printActivityAnalysis(printActivityAnalysis) {}
717766

718767
LogicalResult matchAndRewrite(AutoDiffOp op,
719768
PatternRewriter &rewriter) const override {
@@ -732,12 +781,17 @@ class MinimizeReverseActivity : public OpRewritePattern<AutoDiffOp> {
732781
FailureOr<ActivityOptimization> optimized =
733782
dataflow
734783
? computeDataFlowActivityOptimization(op, callee, argumentActivity,
735-
resultActivity, symbolTable)
784+
resultActivity, symbolTable,
785+
printActivityAnalysis)
736786
: computeClassicalActivityOptimization(op, callee, argumentActivity,
737787
resultActivity, symbolTable);
788+
738789
if (failed(optimized))
739790
return failure();
740791

792+
if (printActivityAnalysis)
793+
printTaggedActivityAnalysis(callee, optimized->taggedActivities);
794+
741795
if (llvm::equal(argumentActivity, optimized->argumentActivity) &&
742796
llvm::equal(resultActivity, optimized->resultActivity))
743797
return failure();
@@ -784,6 +838,7 @@ class MinimizeReverseActivity : public OpRewritePattern<AutoDiffOp> {
784838

785839
private:
786840
bool dataflow;
841+
bool printActivityAnalysis;
787842
};
788843

789844
struct EnzymeActivityOptPass
@@ -793,7 +848,7 @@ struct EnzymeActivityOptPass
793848
void runOnOperation() override {
794849
RewritePatternSet patterns(&getContext());
795850
patterns.add<MinimizeForwardActivity, MinimizeReverseActivity>(
796-
&getContext(), dataflow);
851+
&getContext(), dataflow, printActivityAnalysis);
797852

798853
GreedyRewriteConfig config;
799854
config.enableFolding();

‎enzyme/Enzyme/MLIR/Passes/Passes.td‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,13 @@ def EnzymeActivityOpt : Pass<"enzyme-activity-opt">{
9494
/*default=*/"true",
9595
/*description=*/"Whether to use the new Dataflow activity analysis"
9696
>,
97+
Option<
98+
/*C++ variable name=*/"printActivityAnalysis",
99+
/*CLI argument=*/"print",
100+
/*type=*/"bool",
101+
/*default=*/"false",
102+
/*description=*/"Print the dataflow activity analysis results used by the optimizer"
103+
>,
97104
];
98105
}
99106

‎enzyme/test/MLIR/OptimizeAD/optimize_activity.mlir‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,15 @@
11
// RUN: %eopt --enzyme-activity-opt --split-input-file %s | FileCheck %s
2+
// RUN: %eopt --enzyme-activity-opt=print=true --split-input-file %s 2>&1 | FileCheck %s --check-prefix=PRINT
23

34
module {
45
func.func @square(%x : f64) -> f64{
56
%y = arith.mulf %x, %x : f64
67
return %y : f64
78
}
89

9-
func.func @square2(%x: f32, %y: f32) -> (f32, f32) {
10-
%p = arith.mulf %x, %x : f32
11-
%q = arith.mulf %y, %y : f32
10+
func.func @square2(%x: f32 {enzyme.tag = "x"}, %y: f32 {enzyme.tag = "y"}) -> (f32, f32) {
11+
%p = arith.mulf %x, %x {tag = "p"}: f32
12+
%q = arith.mulf %y, %y {tag = "q"}: f32
1213
return %p, %q : f32, f32
1314
}
1415

@@ -19,12 +20,12 @@ module {
1920
return %p,%q,%dy : f32,f32,f32
2021
}
2122

22-
// Fwddiff but we need to use results of activity analysis
23-
func.func @test2(%x: f32, %y: f32, %dx: f32, %dy : f32) -> (f32,f32,f32) {
24-
%r:3 = enzyme.fwddiff @square2(%x, %dx, %y, %dy) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>] } : (f32,f32,f32,f32) -> (f32,f32,f32)
25-
// CHECK-LABEL: func.func @test2
26-
// CHECK: %{{.*}} = enzyme.fwddiff @square2(%arg0, %arg1, %arg3){{.*}}activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}
27-
28-
return %r#0,%r#1,%r#2 : f32,f32,f32
29-
}
23+
// // Fwddiff but we need to use results of activity analysis
24+
// func.func @test2(%x: f32, %y: f32, %dx: f32, %dy : f32) -> (f32,f32,f32) {
25+
// %r:3 = enzyme.fwddiff @square2(%x, %dx, %y, %dy) { activity=[#enzyme<activity enzyme_dup>, #enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>] } : (f32,f32,f32,f32) -> (f32,f32,f32)
26+
// // CHECK-LABEL: func.func @test2
27+
// // CHECK: %{{.*}} = enzyme.fwddiff @square2(%arg0, %arg1, %arg3){{.*}}activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}ret_activity = [#enzyme<activity enzyme_const>, #enzyme<activity enzyme_dup>]{{.*}}
28+
//
29+
// return %r#0,%r#1,%r#2 : f32,f32,f32
30+
// }
3031
}

0 commit comments

Comments
 (0)