[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)
-			fprintf (acfg->compiled_methods_outfile, "%x\n", method->token);
+		if (!mono_method_is_generic_impl (method) && method->token != 0) {
+			fprintf (acfg->compiled_methods_outfile, "%d\n", method->token);
+		}
 	}
 }
 
@@ -14830,6 +14831,8 @@ aot_assembly (MonoAssembly *ass, guint32 jit_opts, MonoAotOptions *aot_options)
 
 	if (acfg->aot_opts.compiled_methods_outfile && acfg->dedup_phase != DEDUP_COLLECT) {
 		acfg->compiled_methods_outfile = fopen (acfg->aot_opts.compiled_methods_outfile, "w+");
+		fprintf(acfg->compiled_methods_outfile, "%s\n", ass->aname.name);
+
 		if (!acfg->compiled_methods_outfile)
 			aot_printerrf (acfg, "Unable to open compiled-methods-outfile specified file %s\n", acfg->aot_opts.compiled_methods_outfile);
 	}

src/mono/sample/HelloWorld/HelloWorld.csproj

@@ -25,6 +25,8 @@
       LibraryFormat="$(_AotLibraryFormat)"
       Assemblies="@(AotInputAssemblies)"
       OutputDir="$(PublishDir)"
+      CollectCompiledMethods="$(StripILCode)"
+      CompiledMethodsOutputPath="$(MethodTokenFilePath)"
       IntermediateOutputPath="$(IntermediateOutputPath)"
       UseAotDataFile="$(UseAotDataFile)"
       CacheFilePath="$(IntermediateOutputPath)aot_compiler_cache.json"
@@ -35,4 +37,23 @@
       <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>
+
+    <ItemGroup>
+      <MethodTokenFiles Include="$(MethodTokenFilePath)/*.txt" />
+    </ItemGroup>
+
+    <ILStrip
+      TrimIndividualMethods="$(TrimIndividualMethods)"
+      MethodTokenFiles="@(MethodTokenFiles)"
+      AssemblyPath="$(PublishDir)">
+    </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:MethodTokenFilePath=$(MethodTokenFilePath) \
 	'/p:NetTracePath="$(NET_TRACE_PATH)"' \
 	'/p:PgoBinaryPath="$(PGO_BINARY_PATH)"' \
 	'/p:MibcProfilePath="$(MIBC_PROFILE_PATH)"'

src/tasks/AotCompilerTask/MonoAOTCompiler.cs

@@ -152,6 +152,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? CompiledMethodsOutputPath { get; set; }
+
     /// <summary>
     /// File to use for profile-guided optimization, *only* the methods described in the file will be AOT compiled.
     /// </summary>
@@ -711,6 +721,26 @@ private PrecompileArguments GetPrecompileArgumentsFor(ITaskItem assemblyItem, st
             aotArgs.Add("dedup-skip");
         }
 
+        if (CollectCompiledMethods)
+        {
+            if (string.IsNullOrEmpty(CompiledMethodsOutputPath))
+            {
+                Log.LogMessage(MessageImportance.Low, "Skipping collecting the list of aot compiled methods, cause the value of CompiledMethodsOutputPath is empty.");
+            }
+            else
+            {
+                if (!Directory.Exists(CompiledMethodsOutputPath))
+                {
+                    Directory.CreateDirectory(CompiledMethodsOutputPath);
+                }
+                string assemblyFileName = Path.GetFileName(assembly);
+                string assemblyName = assemblyFileName.Replace(".", "_");
+                string outputFileName = assemblyName + "_compiled_methods.txt";
+                string outputFilePath = Path.Combine(CompiledMethodsOutputPath, outputFileName);
+                aotArgs.Add($"compiled-methods-outfile={outputFilePath}");
+            }
+        }
+
         // compute output mode and file names
         if (parsedAotMode == MonoAotMode.LLVMOnly || parsedAotMode == MonoAotMode.LLVMOnlyInterp)
         {

src/tasks/MonoTargetsTasks/ILStrip/ILStrip.cs

@@ -12,45 +12,95 @@
 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;
 
 public class ILStrip : Microsoft.Build.Utilities.Task
 {
     /// <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; }
+
+    /// <summary>
+    /// Methods to be trimmed, identified by method token
+    /// </summary>
+    public ITaskItem[] MethodTokenFiles { get; set; } = Array.Empty<ITaskItem>();
+
+    /// <summary>
+    /// Directory where all the assemblies are stored
+    /// </summary>
+    public string AssemblyPath { get; set; } = "";
+
     public override bool Execute()
     {
-        if (Assemblies.Length == 0)
+        if (!TrimIndividualMethods)
         {
-            throw new ArgumentException($"'{nameof(Assemblies)}' is required.", nameof(Assemblies));
-        }
+            if (Assemblies.Length == 0)
+            {
+                throw new ArgumentException($"'{nameof(Assemblies)}' is required.", nameof(Assemblies));
+            }
 
-        int allowedParallelism = DisableParallelStripping ? 1 : Math.Min(Assemblies.Length, Environment.ProcessorCount);
-        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();
-                                                     });
-
-        if (!result.IsCompleted && !Log.HasLoggedErrors)
-        {
-            Log.LogError("Unknown failure occurred while IL stripping assemblies. Check logs to get more details.");
+            int allowedParallelism = DisableParallelStripping ? 1 : Math.Min(Assemblies.Length, Environment.ProcessorCount);
+            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();
+                                                        });
+
+            if (!result.IsCompleted && !Log.HasLoggedErrors)
+            {
+                Log.LogError("Unknown failure occurred while IL stripping assemblies. Check logs to get more details.");
+            }
+
+            return !Log.HasLoggedErrors;
         }
+        else
+        {
+            if (MethodTokenFiles.Length == 0)
+            {
+                throw new ArgumentException($"'{nameof(MethodTokenFiles)}' is required.", nameof(MethodTokenFiles));
+            }
+
+            if (!Directory.Exists(AssemblyPath))
+            {
+                throw new ArgumentException($"'{nameof(AssemblyPath)}' needs to be a valid path.", nameof(AssemblyPath));
+            }
 
-        return !Log.HasLoggedErrors;
+            int allowedParallelism = DisableParallelStripping ? 1 : Math.Min(MethodTokenFiles.Length, Environment.ProcessorCount);
+            if (BuildEngine is IBuildEngine9 be9)
+                allowedParallelism = be9.RequestCores(allowedParallelism);
+            ParallelLoopResult result = Parallel.ForEach(MethodTokenFiles,
+                                                        new ParallelOptions { MaxDegreeOfParallelism = allowedParallelism },
+                                                        (methodTokenFileItem, state) =>
+                                                        {
+                                                            if (!TrimMethods(methodTokenFileItem, AssemblyPath))
+                                                                state.Stop();
+                                                        });
+
+            if (!result.IsCompleted && !Log.HasLoggedErrors)
+            {
+                Log.LogError("Unknown failure occurred while IL stripping assemblies. Check logs to get more details.");
+            }
+
+            return !Log.HasLoggedErrors;
+        }
     }
 
     private bool StripAssembly(ITaskItem assemblyItem)
@@ -81,4 +131,99 @@ private bool StripAssembly(ITaskItem assemblyItem)
         return true;
     }
 
+    private bool TrimMethods(ITaskItem methodTokenFileItem, string assemblyPath)
+    {
+        string methodTokenFile = methodTokenFileItem.ItemSpec;
+        if (!File.Exists(methodTokenFile))
+        {
+            Log.LogMessage(MessageImportance.Low, $"[ILStrip] {methodTokenFile} doesn't exit.");
+            return true;
+        }
+        string[] log = File.ReadAllLines(methodTokenFile);
+        if (log.Length <= 1)
+        {
+            // Frist line is assembly name
+            Log.LogMessage(MessageImportance.Low, $"[ILStrip] {methodTokenFile} doesn't contain any compiled method token.");
+            return true;
+        }
+        string assemblyName = log[0];
+        string assemblyFilePath = Path.Combine(assemblyPath, (assemblyName + ".dll"));
+        string trimmedAssemblyFilePath = Path.Combine(assemblyPath, (assemblyName + "_new.dll"));
+        if (!File.Exists(assemblyFilePath))
+        {
+            Log.LogMessage(MessageImportance.Low, $"[ILStrip] {assemblyFilePath} doesn't exit.");
+            return true;
+        }
+
+        using (FileStream fs = File.Open(assemblyFilePath, FileMode.Open),
+                os = File.Open(trimmedAssemblyFilePath, FileMode.Create))
+        {
+            MemoryStream memStream = new MemoryStream((int)fs.Length);
+            fs.CopyTo(memStream);
+
+            fs.Position = 0;
+            PEReader peReader = new PEReader(fs, PEStreamOptions.LeaveOpen);
+            MetadataReader mr = peReader.GetMetadataReader();
+
+            Dictionary<int, int> token_to_rva = new Dictionary<int, int>();
+            Dictionary<int, int> method_body_uses = new Dictionary<int, int>();
+
+            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 count))
+                {
+                    method_body_uses[rva]++;
+                }
+                else
+                {
+                    method_body_uses.Add(rva, 1);
+                }
+            }
+
+            for (int i = 1; i < log.Length; i++)
+            {
+                int methodToken2Trim = Convert.ToInt32(log[i]);
+                int rva2Trim = token_to_rva[methodToken2Trim];
+                method_body_uses[rva2Trim]--;
+            }
+
+            foreach (var kvp in method_body_uses)
+            {
+                int rva = kvp.Key;
+                int count = kvp.Value;
+                if (count == 0)
+                {
+                    MethodBodyBlock mb = peReader.GetMethodBody(rva);
+                    int methodSize = mb.Size;
+                    int sectionIndex = peReader.PEHeaders.GetContainingSectionIndex(rva);
+                    int relativeOffset = rva - peReader.PEHeaders.SectionHeaders[sectionIndex].VirtualAddress;
+                    int actualLoc = peReader.PEHeaders.SectionHeaders[sectionIndex].PointerToRawData + relativeOffset;
+
+                    byte[] zeroBuffer = new byte[methodSize];
+                    for (int i = 0; i < methodSize; i++)
+                    {
+                        zeroBuffer[i] = 0x0;
+                    }
+
+                    memStream.Position = actualLoc;
+                    int firstbyte = memStream.ReadByte();
+                    int headerFlag = firstbyte & 0b11;
+                    int headerSize = headerFlag == 2 ? 1 : 4;
+
+                    memStream.Position = actualLoc + headerSize;
+                    memStream.Write(zeroBuffer, 0, methodSize - headerSize);
+                }
+            }
+            memStream.Position = 0;
+            memStream.CopyTo(os);
+        }
+
+        return true;
+    }
 }