Getting Metal shader precompiling working.

* Adds GrMtlPipelineStateBuilder::PrecompileShaders.
* Adds GrMtlPrecompileShaderLibrary.
* Stores metadata with SkSL shaders so we can convert to MSL properly

Change-Id: Id2c23a54cc04ca3b61a639c1cbc7234b697c41ab
Reviewed-on: https://skia-review.googlesource.com/c/skia/+/376856
Reviewed-by: Brian Osman <brianosman@google.com>
Commit-Queue: Jim Van Verth <jvanverth@google.com>
diff --git a/src/gpu/GrPersistentCacheUtils.h b/src/gpu/GrPersistentCacheUtils.h
index 37bf52f..e0d3cbf 100644
--- a/src/gpu/GrPersistentCacheUtils.h
+++ b/src/gpu/GrPersistentCacheUtils.h
@@ -66,7 +66,7 @@
     return writer.snapshotAsData();
 }
 
-static SkFourByteTag GetType(SkReadBuffer* reader) {
+static inline SkFourByteTag GetType(SkReadBuffer* reader) {
     constexpr SkFourByteTag kInvalidTag = ~0;
     int version           = reader->readInt();
     SkFourByteTag typeTag = reader->readUInt();
diff --git a/src/gpu/mtl/GrMtlGpu.h b/src/gpu/mtl/GrMtlGpu.h
index 2bd9b93..221273a 100644
--- a/src/gpu/mtl/GrMtlGpu.h
+++ b/src/gpu/mtl/GrMtlGpu.h
@@ -109,6 +109,8 @@
         this->didWriteToSurface(surface, origin, bounds);
     }
 
+    bool precompileShader(const SkData& key, const SkData& data) override;
+
 private:
     GrMtlGpu(GrDirectContext*, const GrContextOptions&, id<MTLDevice>,
              id<MTLCommandQueue>, GrMTLHandle binaryArchive, MTLFeatureSet);
diff --git a/src/gpu/mtl/GrMtlGpu.mm b/src/gpu/mtl/GrMtlGpu.mm
index 8ac593b..48ad7a8 100644
--- a/src/gpu/mtl/GrMtlGpu.mm
+++ b/src/gpu/mtl/GrMtlGpu.mm
@@ -18,6 +18,7 @@
 #include "src/gpu/mtl/GrMtlBuffer.h"
 #include "src/gpu/mtl/GrMtlCommandBuffer.h"
 #include "src/gpu/mtl/GrMtlOpsRenderPass.h"
+#include "src/gpu/mtl/GrMtlPipelineStateBuilder.h"
 #include "src/gpu/mtl/GrMtlSemaphore.h"
 #include "src/gpu/mtl/GrMtlTexture.h"
 #include "src/gpu/mtl/GrMtlTextureRenderTarget.h"
@@ -1446,6 +1447,10 @@
     cmdEncoder.label = @"resolveTexture";
 }
 
+bool GrMtlGpu::precompileShader(const SkData& key, const SkData& data) {
+    return GrMtlPipelineStateBuilder::PrecompileShaders(this, data);
+}
+
 #if GR_TEST_UTILS
 void GrMtlGpu::testingOnly_startCapture() {
     if (@available(macOS 10.13, iOS 11.0, *)) {
diff --git a/src/gpu/mtl/GrMtlPipelineStateBuilder.h b/src/gpu/mtl/GrMtlPipelineStateBuilder.h
index ee44de4..e4b5179 100644
--- a/src/gpu/mtl/GrMtlPipelineStateBuilder.h
+++ b/src/gpu/mtl/GrMtlPipelineStateBuilder.h
@@ -37,6 +37,8 @@
                                                    const GrProgramDesc&,
                                                    const GrProgramInfo&);
 
+    static bool PrecompileShaders(GrMtlGpu*, const SkData&);
+
 private:
     GrMtlPipelineStateBuilder(GrMtlGpu*, GrRenderTarget*,
                               const GrProgramDesc&, const GrProgramInfo&);
@@ -55,7 +57,7 @@
                                            SkSL::Program::Inputs inputs,
                                            GrContextOptions::ShaderErrorHandler* errorHandler);
     void storeShadersInCache(const SkSL::String shaders[], const SkSL::Program::Inputs inputs[],
-                             bool isSkSL);
+                             SkSL::Program::Settings*, bool isSkSL);
 
     GrGLSLUniformHandler* uniformHandler() override { return &fUniformHandler; }
     const GrGLSLUniformHandler* uniformHandler() const override { return &fUniformHandler; }
diff --git a/src/gpu/mtl/GrMtlPipelineStateBuilder.mm b/src/gpu/mtl/GrMtlPipelineStateBuilder.mm
index 7ee72f5..4de002a 100644
--- a/src/gpu/mtl/GrMtlPipelineStateBuilder.mm
+++ b/src/gpu/mtl/GrMtlPipelineStateBuilder.mm
@@ -72,15 +72,24 @@
 
 void GrMtlPipelineStateBuilder::storeShadersInCache(const SkSL::String shaders[],
                                                     const SkSL::Program::Inputs inputs[],
+                                                    SkSL::Program::Settings* settings,
                                                     bool isSkSL) {
     // Here we shear off the Mtl-specific portion of the Desc in order to create the
     // persistent key. This is because Mtl only caches the MSL code, not the fully compiled
     // program, and that only depends on the base GrProgramDesc data.
     sk_sp<SkData> key = SkData::MakeWithoutCopy(this->desc().asKey(),
                                                 this->desc().initialKeyLength());
-    sk_sp<SkData> data = GrPersistentCacheUtils::PackCachedShaders(isSkSL ? kSKSL_Tag : kMSL_Tag,
-                                                                   shaders,
-                                                                   inputs, kGrShaderTypeCount);
+    sk_sp<SkData> data;
+    if (isSkSL) {
+        // source cache, plus metadata to allow for a complete precompile
+        GrPersistentCacheUtils::ShaderMetadata meta;
+        meta.fSettings = settings;
+        data = GrPersistentCacheUtils::PackCachedShaders(kSKSL_Tag, shaders, inputs,
+                                                         kGrShaderTypeCount, &meta);
+    } else {
+        data = GrPersistentCacheUtils::PackCachedShaders(kMSL_Tag, shaders, inputs,
+                                                         kGrShaderTypeCount);
+    }
     fGpu->getContext()->priv().getPersistentCache()->store(*key, *data);
 }
 
@@ -461,9 +470,9 @@
                 SkSL::String sksl[kGrShaderTypeCount];
                 sksl[kVertex_GrShaderType] = GrShaderUtils::PrettyPrint(fVS.fCompilerString);
                 sksl[kFragment_GrShaderType] = GrShaderUtils::PrettyPrint(fFS.fCompilerString);
-                this->storeShadersInCache(sksl, inputs, true);
+                this->storeShadersInCache(sksl, inputs, &settings, true);
             } else {
-                this->storeShadersInCache(msl, inputs, false);
+                this->storeShadersInCache(msl, inputs, nullptr, false);
             }
         }
     }
@@ -574,3 +583,60 @@
 
 //////////////////////////////////////////////////////////////////////////////
 
+bool GrMtlPipelineStateBuilder::PrecompileShaders(GrMtlGpu* gpu, const SkData& cachedData) {
+    SkReadBuffer reader(cachedData.data(), cachedData.size());
+    SkFourByteTag shaderType = GrPersistentCacheUtils::GetType(&reader);
+
+    auto errorHandler = gpu->getContext()->priv().getShaderErrorHandler();
+
+    SkSL::Program::Settings settings;
+    settings.fSharpenTextures = gpu->getContext()->priv().options().fSharpenMipmappedTextures;
+    GrPersistentCacheUtils::ShaderMetadata meta;
+    meta.fSettings = &settings;
+
+    SkSL::String shaders[kGrShaderTypeCount];
+    SkSL::Program::Inputs inputs[kGrShaderTypeCount];
+    if (!GrPersistentCacheUtils::UnpackCachedShaders(&reader, shaders, inputs, kGrShaderTypeCount,
+                                                     &meta)) {
+        return false;
+    }
+
+    switch (shaderType) {
+        case kMSL_Tag: {
+            GrPrecompileMtlShaderLibrary(gpu, shaders[kVertex_GrShaderType]);
+            GrPrecompileMtlShaderLibrary(gpu, shaders[kFragment_GrShaderType]);
+            break;
+        }
+
+        case kSKSL_Tag: {
+            SkSL::String mslShaders[kGrShaderTypeCount];
+            if (!GrSkSLToMSL(gpu,
+                           shaders[kVertex_GrShaderType],
+                           SkSL::ProgramKind::kVertex,
+                           settings,
+                           &mslShaders[kVertex_GrShaderType],
+                           &inputs[kVertex_GrShaderType],
+                           errorHandler)) {
+                return false;
+            }
+            if (!GrSkSLToMSL(gpu,
+                           shaders[kFragment_GrShaderType],
+                           SkSL::ProgramKind::kFragment,
+                           settings,
+                           &mslShaders[kFragment_GrShaderType],
+                           &inputs[kFragment_GrShaderType],
+                           errorHandler)) {
+                return false;
+            }
+            GrPrecompileMtlShaderLibrary(gpu, mslShaders[kVertex_GrShaderType]);
+            GrPrecompileMtlShaderLibrary(gpu, mslShaders[kFragment_GrShaderType]);
+            break;
+        }
+
+        default: {
+            return false;
+        }
+    }
+
+    return true;
+}
diff --git a/src/gpu/mtl/GrMtlUtil.h b/src/gpu/mtl/GrMtlUtil.h
index c5d4c29..fd92ebe 100644
--- a/src/gpu/mtl/GrMtlUtil.h
+++ b/src/gpu/mtl/GrMtlUtil.h
@@ -89,6 +89,13 @@
                                          GrContextOptions::ShaderErrorHandler* errorHandler);
 
 /**
+ * Attempts to compile an MSL shader asynchronously. We are not concerned about the result, which
+ * will be cached in the Apple shader cache.
+ */
+void GrPrecompileMtlShaderLibrary(const GrMtlGpu* gpu,
+                                  const SkSL::String& msl);
+
+/**
  * Replacement for newLibraryWithSource:options:error that has a timeout.
  */
 id<MTLLibrary> GrMtlNewLibraryWithSource(id<MTLDevice>, NSString* mslCode,
diff --git a/src/gpu/mtl/GrMtlUtil.mm b/src/gpu/mtl/GrMtlUtil.mm
index 859c48e..7b7bae8 100644
--- a/src/gpu/mtl/GrMtlUtil.mm
+++ b/src/gpu/mtl/GrMtlUtil.mm
@@ -113,6 +113,21 @@
     return compiledLibrary;
 }
 
+void GrPrecompileMtlShaderLibrary(const GrMtlGpu* gpu,
+                                  const SkSL::String& msl) {
+    auto nsSource = [[NSString alloc] initWithBytesNoCopy:const_cast<char*>(msl.c_str())
+                                                   length:msl.size()
+                                                 encoding:NSUTF8StringEncoding
+                                             freeWhenDone:NO];
+    // Do nothing after completion for now.
+    // TODO: cache the result somewhere so we can use it later.
+    MTLNewLibraryCompletionHandler completionHandler =
+            ^(id<MTLLibrary> library, NSError* error) {};
+    [gpu->device() newLibraryWithSource:nsSource
+                                options:nil
+                      completionHandler:completionHandler];
+}
+
 // Wrapper to get atomic assignment for compiles and pipeline creation
 class MtlCompileResult : public SkRefCnt {
 public: