Skip to content

Commit

Permalink
BEGIN_PUBLIC
Browse files Browse the repository at this point in the history
Basic wiring for optimizer tools which support baseline profile rewriting.
END_PUBLIC

Support rewriting of profiles in optimizer tools.

This implements the blaze portion of Phase 1 of what is described here:
https://docs.google.com/document/d/14ZXSjMK8V1slSCxigGp7B9tCRO_YeXXdNSNSeDzXSxU/edit?usp=sharing

RELNOTES: Support for optimizers rewriting baseline profiles.
PiperOrigin-RevId: 547910576
Change-Id: I4a66a045f9a34e8fbd6484426286a6005b6c9fe8
  • Loading branch information
Googler authored and copybara-github committed Jul 13, 2023
1 parent da23370 commit 5391486
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 21 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,27 @@ public void registerMigrationRuleError(RuleContext ruleContext) throws RuleError
/* Bazel does not currently support baseline profiles in the final apk. */
@Override
public Artifact getArtProfileForApk(
RuleContext ruleContext, Artifact finalClassesDex, Artifact proguardOutputMap) {
RuleContext ruleContext,
Artifact finalClassesDex,
Artifact proguardOutputMap,
String baselineProfileDir) {
return null;
}

/* Bazel does not currently support baseline profiles in the final apk. */
@Override
public Artifact compileBaselineProfile(
RuleContext ruleContext,
Artifact finalClassesDex,
Artifact proguardOutputMap,
Artifact mergedStaticProfile,
String baselineProfileDir) {
return null;
}

/* Bazel does not currently support baseline profiles in the final apk. */
@Override
public Artifact mergeBaselineProfiles(RuleContext ruleContext, String baselineProfileDir) {
return null;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -531,6 +531,17 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
}
}

BaselineProfileProvider baselineprofileProvider =
ruleContext.getPrerequisite("application_resources", BaselineProfileProvider.PROVIDER);

Artifact baselineProfile = null;
String baselineProfileDir = ruleContext.getLabel().getName() + "-baseline-profile/";
if (baselineprofileProvider == null
&& Allowlist.hasAllowlist(ruleContext, "allow_baseline_profiles_optimizer_integration")
&& Allowlist.isAvailable(ruleContext, "allow_baseline_profiles_optimizer_integration")) {
baselineProfile = androidSemantics.mergeBaselineProfiles(ruleContext, baselineProfileDir);
}

ProguardOutput proguardOutput =
applyProguard(
ruleContext,
Expand All @@ -539,7 +550,9 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
binaryJar,
proguardSpecs,
proguardMapping,
proguardOutputMap);
proguardOutputMap,
baselineProfile,
baselineProfileDir);

// Determine the outputs for the proguard map transition through post-processing and adding of
// desugared library map.
Expand Down Expand Up @@ -682,13 +695,28 @@ public static RuleConfiguredTargetBuilder createAndroidBinary(
proguardOutput.addAllToSet(filesBuilder, finalProguardOutputMap);
}

BaselineProfileProvider baselineprofileProvider =
ruleContext.getPrerequisite("application_resources", BaselineProfileProvider.PROVIDER);
Artifact artProfileZip =
(baselineprofileProvider != null)
? baselineprofileProvider.getArtProfileZip()
: androidSemantics.getArtProfileForApk(
ruleContext, finalClassesDex, finalProguardOutputMap);
Artifact artProfileZip = null;
if (baselineprofileProvider != null) {
// This happens when baseline profiles are provided via starlark.
artProfileZip = baselineprofileProvider.getArtProfileZip();
} else if (baselineProfile == null) {
// This happens when optimizer profile rewriting isn't enabled.
artProfileZip =
androidSemantics.getArtProfileForApk(
ruleContext, finalClassesDex, finalProguardOutputMap, baselineProfileDir);
} else {
// This happens when optimizer profile rewriting is enabled.
artProfileZip =
androidSemantics.compileBaselineProfile(
ruleContext,
finalClassesDex,
// Minified symbols are emitted when rewriting, so map isn't needed.
/* proguardOutputMap= */ null,
proguardOutput.getBaselineProfileOut() == null
? baselineProfile
: proguardOutput.getBaselineProfileOut(),
baselineProfileDir);
}

Artifact unsignedApk =
ruleContext.getImplicitOutputArtifact(AndroidRuleClasses.ANDROID_BINARY_UNSIGNED_APK);
Expand Down Expand Up @@ -1054,7 +1082,9 @@ private static ProguardOutput applyProguard(
Artifact deployJarArtifact,
ImmutableList<Artifact> proguardSpecs,
Artifact proguardMapping,
@Nullable Artifact proguardOutputMap)
@Nullable Artifact proguardOutputMap,
@Nullable Artifact baselineProfile,
String baselineProfileDir)
throws InterruptedException {
Artifact proguardOutputJar =
ruleContext.getImplicitOutputArtifact(AndroidRuleClasses.ANDROID_BINARY_PROGUARD_JAR);
Expand Down Expand Up @@ -1101,7 +1131,9 @@ private static ProguardOutput applyProguard(
proguardOutputJar,
javaSemantics,
getProguardOptimizationPasses(ruleContext),
proguardOutputMap);
proguardOutputMap,
baselineProfile,
baselineProfileDir);
}

@Nullable
Expand Down Expand Up @@ -1131,7 +1163,8 @@ private static ProguardOutput createEmptyProguardAction(
ruleContext,
semantics,
proguardOutputMap,
null);
/* libraryJar= */ null,
/* baselineProfileOutput= */ null);
outputs.addAllToSet(failures);
ruleContext.registerAction(
new FailAction(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,21 @@ default void validateAndroidLibraryRuleContext(RuleContext ruleContext)

/** The artifact for ART profile information. */
Artifact getArtProfileForApk(
RuleContext ruleContext, Artifact finalClassesDex, Artifact proguardOutputMap);
RuleContext ruleContext,
Artifact finalClassesDex,
Artifact proguardOutputMap,
String baselineProfileDir);

/** The merged baseline profiles from the {@code baseline_profiles} attribute. */
Artifact mergeBaselineProfiles(RuleContext ruleContext, String baselineProfileDir);

/** The artifact for ART profile information, given a particular merged profile. */
Artifact compileBaselineProfile(
RuleContext ruleContext,
Artifact finalClassesDex,
Artifact proguardOutputMap,
Artifact mergedStaticProfile,
String baselineProfileDir);

boolean postprocessClassesRewritesMap(RuleContext ruleContext);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ public static final class ProguardOutput {
@Nullable private final Artifact constantStringObfuscatedMapping;
@Nullable private final Artifact libraryJar;
private final Artifact config;
@Nullable private final Artifact baselineProfileOut;

public ProguardOutput(
Artifact outputJar,
Expand All @@ -72,7 +73,8 @@ public ProguardOutput(
@Nullable Artifact usage,
@Nullable Artifact constantStringObfuscatedMapping,
@Nullable Artifact libraryJar,
Artifact config) {
Artifact config,
@Nullable Artifact baselineProfileOut) {
this.outputJar = checkNotNull(outputJar);
this.mapping = mapping;
this.protoMapping = protoMapping;
Expand All @@ -81,10 +83,11 @@ public ProguardOutput(
this.constantStringObfuscatedMapping = constantStringObfuscatedMapping;
this.libraryJar = libraryJar;
this.config = config;
this.baselineProfileOut = baselineProfileOut;
}

public static ProguardOutput createEmpty(Artifact outputJar) {
return new ProguardOutput(outputJar, null, null, null, null, null, null, null);
return new ProguardOutput(outputJar, null, null, null, null, null, null, null, null);
}

public Artifact getOutputJar() {
Expand Down Expand Up @@ -129,6 +132,10 @@ public Artifact getConfig() {
return config;
}

public Artifact getBaselineProfileOut() {
return baselineProfileOut;
}

/** Adds the output artifacts to the given set builder. */
public void addAllToSet(NestedSetBuilder<Artifact> filesBuilder) {
addAllToSet(filesBuilder, null);
Expand Down Expand Up @@ -256,7 +263,8 @@ public static ProguardOutput getProguardOutputs(
RuleContext ruleContext,
JavaSemantics semantics,
@Nullable Artifact proguardOutputMap,
@Nullable Artifact libraryJar)
@Nullable Artifact libraryJar,
@Nullable Artifact baselineProfileOutput)
throws InterruptedException {
boolean mappingRequested = genProguardMapping(ruleContext.attributes());

Expand Down Expand Up @@ -284,7 +292,8 @@ public static ProguardOutput getProguardOutputs(
proguardUsage,
proguardConstantStringMap,
libraryJar,
proguardConfigOutput);
proguardConfigOutput,
baselineProfileOutput);
}

/**
Expand All @@ -304,6 +313,10 @@ public static ProguardOutput getProguardOutputs(
* @param optimizationPasses if not null specifies to break proguard up into multiple passes with
* the given number of optimization passes.
* @param proguardOutputMap mapping generated by Proguard if requested. could be null.
* @param baselineProfileIn Profile to pass the optimizer to perform feedback-directed
* optimization.
* @param baselineProfileDir Directory to write baseline profile artifacts if baseline profile is
* provided.
*/
public static ProguardOutput createOptimizationActions(
RuleContext ruleContext,
Expand All @@ -318,10 +331,11 @@ public static ProguardOutput createOptimizationActions(
Artifact proguardOutputJar,
JavaSemantics semantics,
@Nullable Integer optimizationPasses,
@Nullable Artifact proguardOutputMap)
@Nullable Artifact proguardOutputMap,
@Nullable Artifact baselineProfileIn,
String baselineProfileDir)
throws InterruptedException {
Preconditions.checkArgument(!proguardSpecs.isEmpty());

Artifact libraryJar = null;

if (!libraryJars.isEmpty() && !libraryJars.isSingleton()) {
Expand Down Expand Up @@ -359,7 +373,10 @@ public static ProguardOutput createOptimizationActions(
libraryJars = NestedSetBuilder.create(Order.STABLE_ORDER, filteredLibraryJar);
libraryJar = filteredLibraryJar;
}

Artifact baselineProfileOut = null;
if (baselineProfileIn != null) {
baselineProfileOut = ruleContext.getBinArtifact(baselineProfileDir + "rewritten-prof.txt");
}
ProguardOutput output =
getProguardOutputs(
proguardOutputJar,
Expand All @@ -368,7 +385,8 @@ public static ProguardOutput createOptimizationActions(
ruleContext,
semantics,
proguardOutputMap,
libraryJar);
libraryJar,
baselineProfileOut);

JavaConfiguration javaConfiguration =
ruleContext.getConfiguration().getFragment(JavaConfiguration.class);
Expand All @@ -394,6 +412,8 @@ public static ProguardOutput createOptimizationActions(
output.getUsage(),
output.getConstantStringObfuscatedMapping(),
output.getConfig(),
baselineProfileIn,
baselineProfileOut,
mnemonic);
proguardAction
.setProgressMessage("Trimming binary with %s: %s", mnemonic, ruleContext.getLabel())
Expand Down Expand Up @@ -435,6 +455,8 @@ public static ProguardOutput createOptimizationActions(
/* proguardUsage */ null,
/* constantStringObfuscatedMapping */ null,
/* proguardConfigOutput */ null,
baselineProfileIn,
/* baselineProfileOutput */ null,
mnemonic);
initialAction
.setProgressMessage("Trimming binary with %s: Verification/Shrinking Pass", mnemonic)
Expand Down Expand Up @@ -515,6 +537,8 @@ public static ProguardOutput createOptimizationActions(
output.getUsage(),
output.getConstantStringObfuscatedMapping(),
output.getConfig(),
/* baselineProfileInput */ null,
baselineProfileOut,
mnemonic);
finalAction
.setProgressMessage(
Expand Down Expand Up @@ -576,6 +600,8 @@ private static Artifact createSingleOptimizationAction(
/* proguardUsage */ null,
/* constantStringObfuscatedMapping */ null,
/* proguardConfigOutput */ null,
/* baselineProfileInput */ null,
/* baselineProfileOutput */ null,
mnemonic);
optimizationAction
.setProgressMessage(
Expand Down Expand Up @@ -609,6 +635,8 @@ private static void defaultAction(
@Nullable Artifact proguardUsage,
@Nullable Artifact constantStringObfuscatedMapping,
@Nullable Artifact proguardConfigOutput,
@Nullable Artifact baselineProfileInput,
@Nullable Artifact baselineProfileOutput,
String mnemonic) {

builder
Expand Down Expand Up @@ -680,6 +708,14 @@ private static void defaultAction(
builder.addOutput(proguardConfigOutput);
commandLine.addExecPath("-printconfiguration", proguardConfigOutput);
}
if (baselineProfileInput != null) {
builder.addInput(baselineProfileInput);
commandLine.addExecPath("-baselineprofile", baselineProfileInput);
}
if (baselineProfileOutput != null) {
builder.addOutput(baselineProfileOutput);
commandLine.addExecPath("-printbaselineprofile", baselineProfileOutput);
}
}

/** Returns an intermediate artifact used to run Proguard. */
Expand Down

0 comments on commit 5391486

Please sign in to comment.