[Mono] Add the capability of trimming IL code of individual methods

src/mono/mono/mini/aot-compiler.c

@@ -9828,8 +9828,9 @@ compile_method (MonoAotCompile *acfg, MonoMethod *method)
  mono_atomic_inc_i32 (&acfg->stats.ccount);
 
  if (acfg->aot_opts.compiled_methods_outfile && acfg->compiled_methods_outfile != NULL) {
- if (!mono_method_is_generic_impl (method) && method->token != 0)
+ if (!mono_method_is_generic_impl (method) && method->token != 0) {
  fprintf (acfg->compiled_methods_outfile, "%x\n", method->token);
+ }
  }
 }
 
@@ -14832,6 +14833,10 @@ aot_assembly (MonoAssembly *ass, guint32 jit_opts, MonoAotOptions *aot_options)
  acfg->compiled_methods_outfile = fopen (acfg->aot_opts.compiled_methods_outfile, "w+");
  if (!acfg->compiled_methods_outfile)
  aot_printerrf (acfg, "Unable to open compiled-methods-outfile specified file %s\n", acfg->aot_opts.compiled_methods_outfile);
+ else {
+ fprintf(acfg->compiled_methods_outfile, "%s\n", ass->image->filename);
+ fprintf(acfg->compiled_methods_outfile, "%s\n", ass->image->guid);
+ }
  }
 
  if (acfg->aot_opts.data_outfile) {

src/mono/sample/HelloWorld/HelloWorld.csproj

@@ -25,6 +25,8 @@
  LibraryFormat="$(_AotLibraryFormat)"
  Assemblies="@(AotInputAssemblies)"
  OutputDir="$(PublishDir)"
+ CollectCompiledMethods="$(StripILCode)"
+ CompiledMethodsOutputDirectory="$(CompiledMethodsOutputDirectory)"
  IntermediateOutputPath="$(IntermediateOutputPath)"
  UseAotDataFile="$(UseAotDataFile)"
  CacheFilePath="$(IntermediateOutputPath)aot_compiler_cache.json"
@@ -35,4 +37,18 @@
  <Output TaskParameter="CompiledAssemblies" ItemName="BundleAssemblies" />
  </MonoAOTCompiler>
  </Target>
+
+ <UsingTask TaskName="ILStrip"
+ AssemblyFile="$(MonoTargetsTasksAssemblyPath)" />
+
+ <Target Name="StripILCode" Condition="'$(StripILCode)' == 'true'" AfterTargets="AOTCompileApp">
+ <PropertyGroup>
+ <TrimIndividualMethods>true</TrimIndividualMethods>
+ </PropertyGroup>
+
+ <ILStrip
+ TrimIndividualMethods="$(TrimIndividualMethods)"
+ Assemblies="@(BundleAssemblies)">
+ </ILStrip>
+ </Target>
 </Project>

src/mono/sample/HelloWorld/Makefile

@@ -6,6 +6,8 @@ MONO_CONFIG?=Debug
 MONO_ARCH?=$(shell . $(TOP)eng/native/init-os-and-arch.sh && echo $${arch})
 TARGET_OS?=$(shell . $(TOP)eng/native/init-os-and-arch.sh && echo $${os})
 AOT?=false
+StripILCode?=false
+MethodTokenFilePath?= #<path-to-a-writable-directory>
 
 #NET_TRACE_PATH=<path-to-trace-of-sample>
 #PGO_BINARY_PATH=<path-to-dotnet-pgo-executable>
@@ -18,6 +20,8 @@ publish:
  -c $(MONO_CONFIG) \
  -r $(TARGET_OS)-$(MONO_ARCH) \
  /p:RunAOTCompilation=$(AOT) \
+ /p:StripILCode=$(StripILCode) \
+ /p:CompiledMethodsOutputDirectory=$(CompiledMethodsOutputDirectory) \
  '/p:NetTracePath="$(NET_TRACE_PATH)"' \
  '/p:PgoBinaryPath="$(PGO_BINARY_PATH)"' \
  '/p:MibcProfilePath="$(MIBC_PROFILE_PATH)"'

src/tasks/AotCompilerTask/MonoAOTCompiler.cs

@@ -67,6 +67,7 @@ public class MonoAOTCompiler : Microsoft.Build.Utilities.Task
  /// - LlvmObjectFile (if using LLVM)
  /// - LlvmBitcodeFile (if using LLVM-only)
  /// - ExportsFile (used in LibraryMode only)
+ /// - MethodTokenFile (when using CollectCompiledMethods=true)
  /// </summary>
  [Output]
  public ITaskItem[]? CompiledAssemblies { get; set; }
@@ -152,6 +153,16 @@ public class MonoAOTCompiler : Microsoft.Build.Utilities.Task
  /// </summary>
  public bool UseDwarfDebug { get; set; }
 
+ /// <summary>
+ /// Instructs the AOT compiler to print the list of aot compiled methods
+ /// </summary>
+ public bool CollectCompiledMethods { get; set; }
+
+ /// <summary>
+ /// Directory to store the aot output when using switch compiled-methods-outfile
+ /// </summary>
+ public string? CompiledMethodsOutputDirectory { get; set; }
+
  /// <summary>
  /// File to use for profile-guided optimization, *only* the methods described in the file will be AOT compiled.
  /// </summary>
@@ -436,6 +447,17 @@ private bool ProcessAndValidateArguments()
  throw new LogAsErrorException($"Could not find {fullPath} to AOT");
  }
 
+ if (CollectCompiledMethods)
+ {
+ if (string.IsNullOrEmpty(CompiledMethodsOutputDirectory))
+ throw new LogAsErrorException($"{nameof(CompiledMethodsOutputDirectory)} is empty. When {nameof(CollectCompiledMethods)} is set to true, the user needs to provide a directory for {nameof(CompiledMethodsOutputDirectory)}.");
+
+ if (!Directory.Exists(CompiledMethodsOutputDirectory))
+ {
+ Directory.CreateDirectory(CompiledMethodsOutputDirectory);
+ }
+ }
+
  return !Log.HasLoggedErrors;
  }
 
@@ -711,6 +733,24 @@ private PrecompileArguments GetPrecompileArgumentsFor(ITaskItem assemblyItem, st
  aotArgs.Add("dedup-skip");
  }
 
+ if (CollectCompiledMethods)
+ {
+ string assemblyFileName = Path.GetFileName(assembly);
+ string assemblyName = assemblyFileName.Replace(".", "_");
+ string outputFileName = assemblyName + "_compiled_methods.txt";
+ string outputFilePath;
+ if (string.IsNullOrEmpty(CompiledMethodsOutputDirectory))
+ {
+ outputFilePath = outputFileName;
+ }
+ else
+ {
+ outputFilePath = Path.Combine(CompiledMethodsOutputDirectory, outputFileName);
+ }
+ aotArgs.Add($"compiled-methods-outfile={outputFilePath}");
+ aotAssembly.SetMetadata("MethodTokenFile", outputFilePath);
+ }
+
  // compute output mode and file names
  if (parsedAotMode == MonoAotMode.LLVMOnly || parsedAotMode == MonoAotMode.LLVMOnlyInterp)
  {

src/tasks/MonoTargetsTasks/ILStrip/ILStrip.cs

@@ -12,21 +12,30 @@
 using CilStrip.Mono.Cecil.Binary;
 using CilStrip.Mono.Cecil.Cil;
 using CilStrip.Mono.Cecil.Metadata;
+using System.Reflection.Metadata;
+using System.Reflection.Metadata.Ecma335;
+using System.Reflection.PortableExecutable;
+using System.Buffers;
 
 public class ILStrip : Microsoft.Build.Utilities.Task
 {
+ [Required]
  /// <summary>
  /// Assemblies to be stripped.
  /// The assemblies will be modified in place if OutputPath metadata is not set.
  /// </summary>
- [Required]
  public ITaskItem[] Assemblies { get; set; } = Array.Empty<ITaskItem>();
 
  /// <summary>
  /// Disable parallel stripping
  /// </summary>
  public bool DisableParallelStripping { get; set; }
 
+ /// <summary>
+ /// Enable the feature of trimming indiviual methods
+ /// </summary>
+ public bool TrimIndividualMethods { get; set; }
+
  public override bool Execute()
  {
  if (Assemblies.Length == 0)
@@ -38,12 +47,20 @@ public override bool Execute()
  if (BuildEngine is IBuildEngine9 be9)
  allowedParallelism = be9.RequestCores(allowedParallelism);
  ParallelLoopResult result = Parallel.ForEach(Assemblies,
- new ParallelOptions { MaxDegreeOfParallelism = allowedParallelism },
- (assemblyItem, state) =>
- {
- if (!StripAssembly(assemblyItem))
- state.Stop();
- });
+ new ParallelOptions { MaxDegreeOfParallelism = allowedParallelism },
+ (assemblyItem, state) =>
+ {
+ if (!TrimIndividualMethods)
+ {
+ if (!StripAssembly(assemblyItem))
+ state.Stop();
+ }
+ else
+ {
+ if (!TrimMethods(assemblyItem))
+ state.Stop();
+ }
+ });
 
  if (!result.IsCompleted && !Log.HasLoggedErrors)
  {
@@ -81,4 +98,191 @@ private bool StripAssembly(ITaskItem assemblyItem)
  return true;
  }
 
+ private bool TrimMethods(ITaskItem assemblyItem)
+ {
+ string methodTokenFile = assemblyItem.GetMetadata("MethodTokenFile");
+ if (string.IsNullOrEmpty(methodTokenFile))
+ {
+ Log.LogError($"Metadata MethodTokenFile of {assemblyItem.ItemSpec} is empty");
+ }
+ if (!File.Exists(methodTokenFile))
+ {
+ Log.LogMessage(MessageImportance.Low, $"{methodTokenFile} doesn't exist.");
+ return true;
+ }
+
+ using StreamReader sr = new(methodTokenFile);
+ {
+ string? assemblyFilePath = sr.ReadLine();
+ if (string.IsNullOrEmpty(assemblyFilePath))
+ {
+ return true;
+ }
+
+ if (!File.Exists(assemblyFilePath))
+ {
+ Log.LogError($"{assemblyFilePath} read from {methodTokenFile} doesn't exist.");
+ return true;
+ }
+
+ string trimmedAssemblyFilePath = ComputeTrimmedAssemblyPath(assemblyFilePath);
+ bool isTrimmed = false;
+
+ using FileStream fs = File.Open(assemblyFilePath, FileMode.Open);
+ {
+ using PEReader peReader = new(fs, PEStreamOptions.LeaveOpen);
+ MetadataReader mr = peReader.GetMetadataReader();
+
+ string guidValue = ComputeGuid(mr);
+ string? expectedGuidValue = sr.ReadLine();
+ if (!string.Equals(guidValue, expectedGuidValue, StringComparison.OrdinalIgnoreCase))
+ {
+ Log.LogError($"[ILStrip] GUID value of {assemblyFilePath} doesn't match the value listed in {methodTokenFile}.");
+ return true;
+ }
+
+ string? line = sr.ReadLine();
+ if (!string.IsNullOrEmpty(line))
+ {
+ isTrimmed = true;
+ Dictionary<int, int> method_body_uses = ComputeMethodBodyUsage(mr, sr, line, methodTokenFile);
+ CreateTrimmedAssembly(peReader, trimmedAssemblyFilePath, fs, method_body_uses);
+ }
+ }
+ if (isTrimmed)
+ {
+ ReplaceAssemblyWithTrimmedOne(assemblyFilePath, trimmedAssemblyFilePath);
+ }
+ }
+
+ return true;
+ }
+
+ private static string ComputeTrimmedAssemblyPath(string assemblyFilePath)
+ {
+ string? assemblyPath = Path.GetDirectoryName(assemblyFilePath);
+ string? assemblyName = Path.GetFileNameWithoutExtension(assemblyFilePath);
+ if (string.IsNullOrEmpty(assemblyPath))
+ {
+ return (assemblyName + "_new.dll");
+ }
+ else
+ {
+ return Path.Combine(assemblyPath, (assemblyName + "_new.dll"));
+ }
+ }
+
+ private static string ComputeGuid(MetadataReader mr)
+ {
+ GuidHandle mvidHandle = mr.GetModuleDefinition().Mvid;
+ Guid mvid = mr.GetGuid(mvidHandle);
+ return mvid.ToString();
+ }
+
+ private Dictionary<int, int> ComputeMethodBodyUsage(MetadataReader mr, StreamReader sr, string? line, string methodTokenFile)
+ {
+ Dictionary<int, int> token_to_rva = new();
+ Dictionary<int, int> method_body_uses = new();
+
+ foreach (MethodDefinitionHandle mdefh in mr.MethodDefinitions)
+ {
+ int methodToken = MetadataTokens.GetToken(mr, mdefh);
+ MethodDefinition mdef = mr.GetMethodDefinition(mdefh);
+ int rva = mdef.RelativeVirtualAddress;
+
+ token_to_rva.Add(methodToken, rva);
+
+ if (method_body_uses.TryGetValue(rva, out var _))
+ {
+ method_body_uses[rva]++;
+ }
+ else
+ {
+ method_body_uses.Add(rva, 1);
+ }
+ }
+
+ do
+ {
+ int methodToken2Trim = Convert.ToInt32(line, 16);
+ if (methodToken2Trim <= 0)
+ {
+ Log.LogError($"Method token: {line} in {methodTokenFile} is not a valid hex value.");
+ }
+ if (token_to_rva.TryGetValue(methodToken2Trim, out int rva2Trim))
+ {
+ method_body_uses[rva2Trim]--;
+ }
+ else
+ {
+ Log.LogError($"Method token: {line} in {methodTokenFile} can't be found within the assembly.");
+ }
+ } while ((line = sr.ReadLine()) != null);
+
+ return method_body_uses;
+ }
+
+ private void CreateTrimmedAssembly(PEReader peReader, string trimmedAssemblyFilePath, FileStream fs, Dictionary<int, int> method_body_uses)
+ {
+ using FileStream os = File.Open(trimmedAssemblyFilePath, FileMode.Create);
+ {
+ fs.Position = 0;
+ MemoryStream memStream = new MemoryStream((int)fs.Length);
+ fs.CopyTo(memStream);
+
+ foreach (var kvp in method_body_uses)
+ {
+ int rva = kvp.Key;
+ int count = kvp.Value;
+ if (count == 0)
+ {
+ int methodSize = ComputeMethodSize(peReader, rva);
+ int actualLoc = ComputeMethodHash(peReader, rva);
+ int headerSize = ComputeMethodHeaderSize(memStream, actualLoc);
+ ZeroOutMethodBody(ref memStream, methodSize, actualLoc, headerSize);
+ }
+ else if (count < 0)
+ {
+ Log.LogError($"Method usage count is less than zero for rva: {rva}.");
+ }
+ }
+
+ memStream.Position = 0;
+ memStream.CopyTo(os);
+ }
+ }
+
+ private static int ComputeMethodSize(PEReader peReader, int rva) => peReader.GetMethodBody(rva).Size;
+
+ private static int ComputeMethodHash(PEReader peReader, int rva)
+ {
+ int sectionIndex = peReader.PEHeaders.GetContainingSectionIndex(rva);
+ int relativeOffset = rva - peReader.PEHeaders.SectionHeaders[sectionIndex].VirtualAddress;
+ return (peReader.PEHeaders.SectionHeaders[sectionIndex].PointerToRawData + relativeOffset);
+ }
+
+ private static int ComputeMethodHeaderSize(MemoryStream memStream, int actualLoc)
+ {
+ memStream.Position = actualLoc;
+ int firstbyte = memStream.ReadByte();
+ int headerFlag = firstbyte & 0b11;
+ return (headerFlag == 2 ? 1 : 4);
+ }
+
+ private static void ZeroOutMethodBody(ref MemoryStream memStream, int methodSize, int actualLoc, int headerSize)
+ {
+ memStream.Position = actualLoc + headerSize;
+
+ byte[] zeroBuffer;
+ zeroBuffer = ArrayPool<byte>.Shared.Rent(methodSize);
+ Array.Clear(zeroBuffer, 0, zeroBuffer.Length);
+ memStream.Write(zeroBuffer, 0, methodSize - headerSize);
+ ArrayPool<byte>.Shared.Return(zeroBuffer);
+ }
+
+ private static void ReplaceAssemblyWithTrimmedOne(string assemblyFilePath, string trimmedAssemblyFilePath)
+ {
+ File.Delete(assemblyFilePath);
+ File.Move(trimmedAssemblyFilePath, assemblyFilePath);
+ }
 }