[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);
+    }
 }