[VFX] Collect CustomHLSL Per Context

Fix specific failure scenario (initially identified by @ludovic-theobald during GDC) which is a consequence of bad fetching of customhlsl expression: Instead of per system, it should be done per context. It has two effects:
- Fix attribute generation: modifying non stored variable was leading to compilation failure.
- Allow diverging Buffer usage: Same buffer used as RWBuffer & Buffer in two separate context.

🎁 Fix untracked issue about multiple include within CustomHLSL (see 813e3a7dff57e3192c35f0dd88a248a48de37572)

Author: PaulDemeulenaere

Packages/com.unity.visualeffectgraph/Editor/Compiler/VFXExpressionGraph.cs

@@ -56,7 +56,7 @@ private void CompileExpressionContext(IEnumerable<VFXContext> contexts,
                 if (mapper != null)
                 {
                     foreach (var exp in mapper.expressions)
-                        expressionContext.RegisterExpression(exp);
+                        expressionContext.RegisterExpression(exp, context);
                     contextsToExpressions.Add(context, mapper);
                 }
             }
@@ -82,22 +82,18 @@ private void CompileExpressionContext(IEnumerable<VFXContext> contexts,
             foreach (var exp in expressionsToReduced.Values)
                 AddExpressionDataRecursively(m_ExpressionsData, exp);
 
-            var bufferUsage = m_BufferUsage
-                .Concat(expressionContext.GraphicsBufferUsageType)
-                .GroupBy(o => o.Key).ToArray();
-
-            m_BufferUsage.Clear();
-            foreach (var expression in bufferUsage)
+            foreach (var bufferTypeUsage in expressionContext.GraphicsBufferTypeUsagePerContext)
             {
-                var types = expression.Select(o => o.Value);
-                if (types.Count() != 1)
-                    throw new InvalidOperationException("Diverging type usage for GraphicsBuffer : " + types.Select(o => o.ToString()).Aggregate((a, b) => a + b));
-                m_BufferUsage.Add(expression.Key, types.First());
+                m_BufferTypeUsagePerContext.TryAdd(bufferTypeUsage.Key, bufferTypeUsage.Value);
             }
 
+
             if (target == VFXDeviceTarget.GPU)
             {
-                m_CustomHLSLExpressions = expressionContext.hlslCodeHolders;
+                foreach (var hlslCodeHolder in expressionContext.hlslCodeHoldersPerContext)
+                {
+                    m_CustomHLSLExpressionsPerContext.Add(hlslCodeHolder.Key, hlslCodeHolder.Value);
+                }
             }
         }
 
@@ -294,12 +290,27 @@ private VFXExpressionMapper BuildMapper(VFXContext context, Dictionary<VFXContex
 
         public IEnumerable<VFXLayoutElementDesc> GlobalEventAttributes => m_GlobalEventAttributes;
 
-        public ReadOnlyDictionary<VFXExpression, BufferUsage> BufferUsage => new ReadOnlyDictionary<VFXExpression, BufferUsage>(m_BufferUsage);
+        public ReadOnlyDictionary<VFXExpression, BufferUsage> GetBufferTypeUsage(VFXContext context)
+        {
+            if (m_BufferTypeUsagePerContext.TryGetValue(context, out var bufferTypeUsage))
+            {
+                return new ReadOnlyDictionary<VFXExpression, BufferUsage>(bufferTypeUsage);
+            }
 
-        public IHLSLCodeHolder[] customHLSLExpressions => m_CustomHLSLExpressions;
+            return new ReadOnlyDictionary<VFXExpression, BufferUsage>(new Dictionary<VFXExpression, BufferUsage>());
+        }
+
+        public IHLSLCodeHolder[] GetCustomHLSLExpressions(VFXContext context)
+        {
+            if (m_CustomHLSLExpressionsPerContext.TryGetValue(context, out var hlslCodeHolders))
+            {
+                return hlslCodeHolders.ToArray();
+            }
+            return Array.Empty<IHLSLCodeHolder>();
+        }
 
-        private IHLSLCodeHolder[] m_CustomHLSLExpressions;
-        private Dictionary<VFXExpression, BufferUsage> m_BufferUsage = new Dictionary<VFXExpression, BufferUsage>();
+        private Dictionary<VFXContext, List<IHLSLCodeHolder>> m_CustomHLSLExpressionsPerContext = new();
+        private Dictionary<VFXContext, Dictionary<VFXExpression, BufferUsage>> m_BufferTypeUsagePerContext = new();
         private HashSet<VFXExpression> m_Expressions = new HashSet<VFXExpression>();
         private Dictionary<VFXExpression, VFXExpression> m_CPUExpressionsToReduced = new Dictionary<VFXExpression, VFXExpression>();
         private Dictionary<VFXExpression, VFXExpression> m_GPUExpressionsToReduced = new Dictionary<VFXExpression, VFXExpression>();

Packages/com.unity.visualeffectgraph/Editor/Compiler/VFXGraphCompiledData.cs

@@ -842,7 +842,7 @@ private void GenerateShaders(List<GeneratedCodeData> outGeneratedCodeData, VFXEx
                         var contextData = compiledData.taskToCompiledData[task];
                         contextData.gpuMapper = gpuMapper;
                         contextData.uniformMapper = uniformMapper;
-                        contextData.bufferUsage = graph.BufferUsage;
+                        contextData.bufferUsage = graph.GetBufferTypeUsage(context);
 
                         if (task.doesGenerateShader)
                         {
@@ -1183,7 +1183,7 @@ public void Compile(VFXCompilationMode compilationMode, bool forceShaderValidati
                     foreach (var task in contextCompiledData.tasks)
                     {
                         var contextData = new VFXTaskCompiledData() { indexInShaderSource = -1 };
-                        contextData.hlslCodeHolders = m_ExpressionGraph.customHLSLExpressions;
+                        contextData.hlslCodeHolders = m_ExpressionGraph.GetCustomHLSLExpressions(context);
                         contextData.cpuMapper = cpuMapper;
                         contextData.parameters = context.additionalMappings.ToArray();
                         contextData.linkedEventOut = ComputeEventListFromSlot(context.allLinkedOutputSlot).ToArray();