Reimplement afdo support for rust

Ignore-AOSP-First: The parent CL is internal
Bug: 267229065
Test: go test
Change-Id: Ia14679285b92f3f14ff269392a61f978c71311b2
Merged-In: Ia14679285b92f3f14ff269392a61f978c71311b2

cc/afdo.go

@@ -24,6 +24,7 @@ import (
 	"github.com/google/blueprint/proptools"
 )
 
+// TODO(b/267229066): Remove globalAfdoProfileProjects after implementing bp2build converter for fdo_profile
 var (
 	globalAfdoProfileProjects = []string{
 		"vendor/google_data/pgo_profile/sampling/",
@@ -35,13 +36,6 @@ var afdoProfileProjectsConfigKey = android.NewOnceKey("AfdoProfileProjects")
 
 const afdoCFlagsFormat = "-funique-internal-linkage-names -fprofile-sample-accurate -fprofile-sample-use=%s"
 
-// TODO(b/267229065): Remove getAfdoProfileProjects after reimplementing afdo support for rust
-func getAfdoProfileProjects(config android.DeviceConfig) []string {
-	return config.OnceStringSlice(afdoProfileProjectsConfigKey, func() []string {
-		return globalAfdoProfileProjects
-	})
-}
-
 func recordMissingAfdoProfileFile(ctx android.BaseModuleContext, missing string) {
 	getNamedMapForConfig(ctx.Config(), modulesMissingProfileFileKey).Store(missing, true)
 }
@@ -75,38 +69,6 @@ func (afdo *afdo) afdoEnabled() bool {
 	return afdo != nil && afdo.Properties.Afdo && afdo.Properties.FdoProfilePath != nil
 }
 
-// Get list of profile file names, ordered by level of specialisation. For example:
-//  1. libfoo_arm64.afdo
-//  2. libfoo.afdo
-//
-// Add more specialisation as needed.
-// TODO(b/267229065): Remove getProfileFiles after reimplementing afdo support for rust
-func getProfileFiles(ctx android.BaseModuleContext, moduleName string) []string {
-	var files []string
-	files = append(files, moduleName+"_"+ctx.Arch().ArchType.String()+".afdo")
-	files = append(files, moduleName+".afdo")
-	return files
-}
-
-// TODO(b/267229065): Remove GetAfdoProfileFile after reimplementing afdo support for rust
-func (props *AfdoProperties) GetAfdoProfileFile(ctx android.BaseModuleContext, module string) android.OptionalPath {
-	// Test if the profile_file is present in any of the Afdo profile projects
-	for _, profileFile := range getProfileFiles(ctx, module) {
-		for _, profileProject := range getAfdoProfileProjects(ctx.DeviceConfig()) {
-			path := android.ExistentPathForSource(ctx, profileProject, profileFile)
-			if path.Valid() {
-				return path
-			}
-		}
-	}
-
-	// Record that this module's profile file is absent
-	missing := ctx.ModuleDir() + ":" + module
-	recordMissingAfdoProfileFile(ctx, missing)
-
-	return android.OptionalPathForPath(nil)
-}
-
 func (afdo *afdo) flags(ctx ModuleContext, flags Flags) Flags {
 	if path := afdo.Properties.FdoProfilePath; path != nil {
 		profileUseFlag := fmt.Sprintf(afdoCFlagsFormat, *path)
@@ -149,7 +111,7 @@ func (afdo *afdo) addDep(ctx BaseModuleContext, actx android.BottomUpMutatorCont
 
 // FdoProfileMutator reads the FdoProfileProvider from a direct dep with FdoProfileTag
 // assigns FdoProfileInfo.Path to the FdoProfilePath mutated property
-func (c *Module) FdoProfileMutator(ctx android.BottomUpMutatorContext) {
+func (c *Module) fdoProfileMutator(ctx android.BottomUpMutatorContext) {
 	if !c.Enabled() {
 		return
 	}

cc/afdo_test.go

@@ -23,11 +23,6 @@ import (
 	"github.com/google/blueprint"
 )
 
-var prepareForTestWithFdoProfile = android.FixtureRegisterWithContext(func(ctx android.RegistrationContext) {
-	ctx.RegisterModuleType("soong_namespace", android.NamespaceFactory)
-	ctx.RegisterModuleType("fdo_profile", fdoProfileFactory)
-})
-
 type visitDirectDepsInterface interface {
 	VisitDirectDeps(blueprint.Module, func(dep blueprint.Module))
 }
@@ -65,7 +60,7 @@ func TestAfdoDeps(t *testing.T) {
 	`
 
 	result := android.GroupFixturePreparers(
-		prepareForTestWithFdoProfile,
+		PrepareForTestWithFdoProfile,
 		prepareForCcTest,
 		android.FixtureAddTextFile("afdo_profiles_package/libTest.afdo", ""),
 		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
@@ -160,7 +155,7 @@ func TestAfdoEnabledOnStaticDepNoAfdo(t *testing.T) {
 
 	result := android.GroupFixturePreparers(
 		prepareForCcTest,
-		prepareForTestWithFdoProfile,
+		PrepareForTestWithFdoProfile,
 		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/libFoo.afdo", ""),
 		android.MockFS{
 			"afdo_profiles_package/Android.bp": []byte(`
@@ -222,7 +217,7 @@ func TestAfdoEnabledWithRuntimeDepNoAfdo(t *testing.T) {
 
 	result := android.GroupFixturePreparers(
 		prepareForCcTest,
-		prepareForTestWithFdoProfile,
+		PrepareForTestWithFdoProfile,
 		android.FixtureAddTextFile("afdo_profiles_package/libTest.afdo", ""),
 		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
 			variables.AfdoProfiles = []string{
@@ -257,7 +252,7 @@ func TestAfdoEnabledWithMultiArchs(t *testing.T) {
 	}
 `
 	result := android.GroupFixturePreparers(
-		prepareForTestWithFdoProfile,
+		PrepareForTestWithFdoProfile,
 		prepareForCcTest,
 		android.FixtureAddTextFile("afdo_profiles_package/foo_arm.afdo", ""),
 		android.FixtureAddTextFile("afdo_profiles_package/foo_arm64.afdo", ""),
@@ -322,7 +317,7 @@ func TestMultipleAfdoRDeps(t *testing.T) {
 	`
 
 	result := android.GroupFixturePreparers(
-		prepareForTestWithFdoProfile,
+		PrepareForTestWithFdoProfile,
 		prepareForCcTest,
 		android.FixtureAddTextFile("afdo_profiles_package/libTest.afdo", ""),
 		android.FixtureAddTextFile("afdo_profiles_package/libBar.afdo", ""),

cc/fdo_profile.go

@@ -50,7 +50,7 @@ var FdoProfileProvider = blueprint.NewMutatorProvider(FdoProfileInfo{}, "fdo_pro
 // module types that can depend on an fdo_profile module
 type FdoProfileMutatorInterface interface {
 	// FdoProfileMutator eithers set or get FdoProfileProvider
-	FdoProfileMutator(ctx android.BottomUpMutatorContext)
+	fdoProfileMutator(ctx android.BottomUpMutatorContext)
 }
 
 var _ FdoProfileMutatorInterface = (*fdoProfile)(nil)
@@ -60,7 +60,7 @@ func (fp *fdoProfile) GenerateAndroidBuildActions(ctx android.ModuleContext) {}
 
 // FdoProfileMutator sets FdoProfileProvider to fdo_profile module
 // or sets afdo.Properties.FdoProfilePath to path in FdoProfileProvider of the depended fdo_profile
-func (fp *fdoProfile) FdoProfileMutator(ctx android.BottomUpMutatorContext) {
+func (fp *fdoProfile) fdoProfileMutator(ctx android.BottomUpMutatorContext) {
 	if fp.properties.Profile != nil {
 		path := android.PathForModuleSrc(ctx, *fp.properties.Profile)
 		ctx.SetProvider(FdoProfileProvider, FdoProfileInfo{
@@ -69,11 +69,11 @@ func (fp *fdoProfile) FdoProfileMutator(ctx android.BottomUpMutatorContext) {
 	}
 }
 
-// fdoProfileMutator calls the generic FdoProfileMutator function of FdoProfileMutator
+// fdoProfileMutator calls the generic fdoProfileMutator function of fdoProfileMutator
 // which is implemented by cc and cc.FdoProfile
 func fdoProfileMutator(ctx android.BottomUpMutatorContext) {
 	if f, ok := ctx.Module().(FdoProfileMutatorInterface); ok {
-		f.FdoProfileMutator(ctx)
+		f.fdoProfileMutator(ctx)
 	}
 }
 

cc/testing.go

@@ -670,6 +670,12 @@ var PrepareForTestWithHostMusl = android.GroupFixturePreparers(
 	`),
 )
 
+// PrepareForTestWithFdoProfile registers module types to test with fdo_profile
+var PrepareForTestWithFdoProfile = android.FixtureRegisterWithContext(func(ctx android.RegistrationContext) {
+	ctx.RegisterModuleType("soong_namespace", android.NamespaceFactory)
+	ctx.RegisterModuleType("fdo_profile", fdoProfileFactory)
+})
+
 // TestConfig is the legacy way of creating a test Config for testing cc modules.
 //
 // See testCc for an explanation as to how to stop using this deprecated method.

rust/afdo.go

@@ -17,7 +17,10 @@ package rust
 import (
 	"fmt"
 
+	"android/soong/android"
 	"android/soong/cc"
+
+	"github.com/google/blueprint"
 )
 
 const afdoFlagFormat = "-Zprofile-sample-use=%s"
@@ -30,19 +33,49 @@ func (afdo *afdo) props() []interface{} {
 	return []interface{}{&afdo.Properties}
 }
 
-func (afdo *afdo) flags(ctx ModuleContext, flags Flags, deps PathDeps) (Flags, PathDeps) {
+func (afdo *afdo) addDep(ctx BaseModuleContext, actx android.BottomUpMutatorContext) {
+	// afdo is not supported outside of Android
+	if ctx.Host() {
+		return
+	}
+
+	if mod, ok := ctx.Module().(*Module); ok && mod.Enabled() {
+		fdoProfileName, err := actx.DeviceConfig().AfdoProfile(actx.ModuleName())
+		if err != nil {
+			ctx.ModuleErrorf("%s", err.Error())
+		}
+		if fdoProfileName != nil {
+			actx.AddFarVariationDependencies(
+				[]blueprint.Variation{
+					{Mutator: "arch", Variation: actx.Target().ArchVariation()},
+					{Mutator: "os", Variation: "android"},
+				},
+				cc.FdoProfileTag,
+				[]string{*fdoProfileName}...,
+			)
+		}
+	}
+}
+
+func (afdo *afdo) flags(ctx android.ModuleContext, flags Flags, deps PathDeps) (Flags, PathDeps) {
 	if ctx.Host() {
 		return flags, deps
 	}
 
-	if afdo != nil && afdo.Properties.Afdo {
-		if profileFile := afdo.Properties.GetAfdoProfileFile(ctx, ctx.ModuleName()); profileFile.Valid() {
-			profileUseFlag := fmt.Sprintf(afdoFlagFormat, profileFile)
+	if !afdo.Properties.Afdo {
+		return flags, deps
+	}
+
+	ctx.VisitDirectDepsWithTag(cc.FdoProfileTag, func(m android.Module) {
+		if ctx.OtherModuleHasProvider(m, cc.FdoProfileProvider) {
+			info := ctx.OtherModuleProvider(m, cc.FdoProfileProvider).(cc.FdoProfileInfo)
+			path := info.Path
+			profileUseFlag := fmt.Sprintf(afdoFlagFormat, path.String())
 			flags.RustFlags = append(flags.RustFlags, profileUseFlag)
 
-			profileFilePath := profileFile.Path()
-			deps.AfdoProfiles = append(deps.AfdoProfiles, profileFilePath)
+			deps.AfdoProfiles = append(deps.AfdoProfiles, path)
 		}
-	}
+	})
+
 	return flags, deps
 }

rust/afdo_test.go

@@ -16,6 +16,7 @@ package rust
 
 import (
 	"android/soong/android"
+	"android/soong/cc"
 	"fmt"
 	"strings"
 	"testing"
@@ -31,13 +32,27 @@ func TestAfdoEnabled(t *testing.T) {
 `
 	result := android.GroupFixturePreparers(
 		prepareForRustTest,
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo.afdo", ""),
+		cc.PrepareForTestWithFdoProfile,
+		android.FixtureAddTextFile("afdo_profiles_package/foo.afdo", ""),
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.AfdoProfiles = []string{
+				"foo://afdo_profiles_package:foo_afdo",
+			}
+		}),
+		android.MockFS{
+			"afdo_profiles_package/Android.bp": []byte(`
+				fdo_profile {
+					name: "foo_afdo",
+					profile: "foo.afdo",
+				}
+			`),
+		}.AddToFixture(),
 		rustMockedFiles.AddToFixture(),
 	).RunTestWithBp(t, bp)
 
 	foo := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
 
-	expectedCFlag := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo.afdo")
+	expectedCFlag := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo.afdo")
 
 	if !strings.Contains(foo.Args["rustcFlags"], expectedCFlag) {
 		t.Errorf("Expected 'foo' to enable afdo, but did not find %q in cflags %q", expectedCFlag, foo.Args["rustcFlags"])
@@ -55,16 +70,37 @@ func TestAfdoEnabledWithMultiArchs(t *testing.T) {
 `
 	result := android.GroupFixturePreparers(
 		prepareForRustTest,
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm.afdo", ""),
-		android.FixtureAddTextFile("toolchain/pgo-profiles/sampling/foo_arm64.afdo", ""),
+		cc.PrepareForTestWithFdoProfile,
+		android.FixtureAddTextFile("afdo_profiles_package/foo_arm.afdo", ""),
+		android.FixtureAddTextFile("afdo_profiles_package/foo_arm64.afdo", ""),
+		android.FixtureModifyProductVariables(func(variables android.FixtureProductVariables) {
+			variables.AfdoProfiles = []string{
+				"foo://afdo_profiles_package:foo_afdo",
+			}
+		}),
+		android.MockFS{
+			"afdo_profiles_package/Android.bp": []byte(`
+				fdo_profile {
+					name: "foo_afdo",
+					arch: {
+						arm: {
+							profile: "foo_arm.afdo",
+						},
+						arm64: {
+							profile: "foo_arm64.afdo",
+						}
+					}
+				}
+			`),
+		}.AddToFixture(),
 		rustMockedFiles.AddToFixture(),
 	).RunTestWithBp(t, bp)
 
 	fooArm := result.ModuleForTests("foo", "android_arm_armv7-a-neon").Rule("rustc")
 	fooArm64 := result.ModuleForTests("foo", "android_arm64_armv8-a").Rule("rustc")
 
-	expectedCFlagArm := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm.afdo")
-	expectedCFlagArm64 := fmt.Sprintf(afdoFlagFormat, "toolchain/pgo-profiles/sampling/foo_arm64.afdo")
+	expectedCFlagArm := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo_arm.afdo")
+	expectedCFlagArm64 := fmt.Sprintf(afdoFlagFormat, "afdo_profiles_package/foo_arm64.afdo")
 
 	if !strings.Contains(fooArm.Args["rustcFlags"], expectedCFlagArm) {
 		t.Errorf("Expected 'fooArm' to enable afdo, but did not find %q in cflags %q", expectedCFlagArm, fooArm.Args["rustcFlags"])

rust/rust.go

@@ -39,7 +39,6 @@ func init() {
 		ctx.BottomUp("rust_libraries", LibraryMutator).Parallel()
 		ctx.BottomUp("rust_stdlinkage", LibstdMutator).Parallel()
 		ctx.BottomUp("rust_begin", BeginMutator).Parallel()
-
 	})
 	android.PostDepsMutators(func(ctx android.RegisterMutatorsContext) {
 		ctx.BottomUp("rust_sanitizers", rustSanitizerRuntimeMutator).Parallel()
@@ -920,7 +919,7 @@ func (mod *Module) GenerateAndroidBuildActions(actx android.ModuleContext) {
 
 	// Calculate rustc flags
 	if mod.afdo != nil {
-		flags, deps = mod.afdo.flags(ctx, flags, deps)
+		flags, deps = mod.afdo.flags(actx, flags, deps)
 	}
 	if mod.compiler != nil {
 		flags = mod.compiler.compilerFlags(ctx, flags)
@@ -1613,6 +1612,8 @@ func (mod *Module) DepsMutator(actx android.BottomUpMutatorContext) {
 
 	// proc_macros are compiler plugins, and so we need the host arch variant as a dependendcy.
 	actx.AddFarVariationDependencies(ctx.Config().BuildOSTarget.Variations(), procMacroDepTag, deps.ProcMacros...)
+
+	mod.afdo.addDep(ctx, actx)
 }
 
 // addRlibDependency will add an rlib dependency, rewriting to the snapshot library if available.