Custom Metal shader uniforms working. (?)

Author: slembcke

CCRendererGLSupport.m

@@ -242,16 +242,4 @@ -(void)bind:(BOOL)bind
 	glBindVertexArray(bind ? _vao : 0);
 }
 
--(void)prepare
-{
-	[_vertexBuffer prepare];
-	[_indexBuffer prepare];
-}
-
--(void)commit
-{
-	[_vertexBuffer commit];
-	[_indexBuffer commit];
-}
-
 @end

cocos2d/CCNoARC.m

@@ -242,7 +242,7 @@ @implementation CCRenderStateGL
 		NSDictionary *globalShaderUniforms = renderer->_globalShaderUniforms;
 		NSDictionary *setters = self->_shader->_uniformSetters;
 		for(NSString *uniformName in setters){
-			CCGLUniformSetter setter = setters[uniformName];
+			CCUniformSetter setter = setters[uniformName];
 			setter(renderer, self->_shaderUniforms, globalShaderUniforms);
 		}
 	}
@@ -288,13 +288,6 @@ -(void)invokeOnRenderer:(CCRenderer *)renderer
 @interface CCRenderStateMetal : CCRenderState @end
 @implementation CCRenderStateMetal {
 	id<MTLRenderPipelineState> _renderPipelineState;
-	
-	NSRange _textureRange;
-	id<MTLSamplerState> _samplers[CCMTL_MAX_TEXTURES];
-	id<MTLTexture> _textures[CCMTL_MAX_TEXTURES];
-	
-	@public
-	BOOL _uniformsPrepared;
 }
 
 // Using GL enums for CCBlendMode types should never have happened. Oops.
@@ -350,40 +343,25 @@ @implementation CCRenderStateMetal {
 
 		self->_renderPipelineState = [[context.device newRenderPipelineStateWithDescriptor:pipelineStateDescriptor error:nil] retain];
 	}
-	
-	if(!self->_uniformsPrepared){
-		CCTexture *mainTexture = self->_shaderUniforms[CCShaderUniformMainTexture];
-		
-		self->_textureRange = NSMakeRange(0, 1);
-		self->_samplers[0] = [mainTexture.metalSampler retain];
-		self->_textures[0] = [mainTexture.metalTexture retain];
-		
-		self->_uniformsPrepared = YES;
-	}
 }
 
 static void
 CCRenderStateMetalTransition(CCRenderStateMetal *self, CCRenderer *renderer, CCRenderStateMetal *previous)
 {
-	CCMetalContext *context = renderer->_metalContext;
+	CCGraphicsBufferBindingsMetal *buffers = (CCGraphicsBufferBindingsMetal *)renderer->_buffers;
+	CCMetalContext *context = buffers->_context;
 	id<MTLRenderCommandEncoder> renderEncoder = context->_currentRenderCommandEncoder;
-	[renderEncoder setRenderPipelineState:self->_renderPipelineState];
 
-//	CCGraphicsBufferMetal *uniformGraphicsBuffer = nil;
-//	id<MTLBuffer> uniformMetalBuffer = uniformGraphicsBuffer->_buffer;
-//	
-//	// Set the global uniform buffers.
-//	[renderEncoder setVertexBuffer:uniformMetalBuffer offset:0 atIndex:1];
-//	[renderEncoder setFragmentBuffer:uniformMetalBuffer offset:0 atIndex:1];
-//	
-//	// Set the uniform buffers.
-//	CCGraphicsBufferPushElements(uniformGraphicsBuffer, # of bytes);
-//	[renderEncoder setVertexBuffer:uniformMetalBuffer offset:offset atIndex:2];
-//	[renderEncoder setFragmentBuffer:uniformMetalBuffer offset:offset atIndex:2];
+	// Bind pipeline state.
+	[renderEncoder setRenderPipelineState:self->_renderPipelineState];
 
-	NSRange range = self->_textureRange;
-	[renderEncoder setFragmentSamplerStates:self->_samplers withRange:range];
-	[renderEncoder setFragmentTextures:self->_textures withRange:range];
+	// Set shader arguments.
+	NSDictionary *globalShaderUniforms = renderer->_globalShaderUniforms;
+	NSDictionary *setters = self->_shader->_uniformSetters;
+	for(NSString *uniformName in setters){
+		CCUniformSetter setter = setters[uniformName];
+		setter(renderer, self->_shaderUniforms, globalShaderUniforms);
+	}
 }
 
 -(void)transitionRenderer:(CCRenderer *)renderer FromState:(CCRenderState *)previous
@@ -411,9 +389,10 @@ -(instancetype)initWithMode:(CCRenderCommandDrawMode)mode renderState:(CCRenderS
 
 -(void)invokeOnRenderer:(CCRenderer *)renderer
 {
-	CCMetalContext *context = renderer->_metalContext;
+	CCGraphicsBufferBindingsMetal *buffers = (CCGraphicsBufferBindingsMetal *)renderer->_buffers;
+	CCMetalContext *context = buffers->_context;
 	id<MTLRenderCommandEncoder> renderEncoder = context->_currentRenderCommandEncoder;
-	id<MTLBuffer> indexBuffer = ((CCGraphicsBufferMetal *)renderer->_elementBuffer)->_buffer;
+	id<MTLBuffer> indexBuffer = ((CCGraphicsBufferMetal *)buffers->_indexBuffer)->_buffer;
 
 	CCMTL_DEBUG_PUSH_GROUP_MARKER(renderEncoder, @"CCRendererCommandDraw: Invoke");
 	CCRendererBindBuffers(renderer, YES);
@@ -423,13 +402,6 @@ -(void)invokeOnRenderer:(CCRenderer *)renderer
 	[renderEncoder drawIndexedPrimitives:MetalDrawModes[_mode] indexCount:_count indexType:MTLIndexTypeUInt16 indexBuffer:indexBuffer indexBufferOffset:2*_first];
 	CCMTL_DEBUG_POP_GROUP_MARKER(renderEncoder);
 
-	if(!_renderState->_immutable){
-		// This is sort of a weird place to put this, but couldn't find somewhere better.
-		// Mutable render states need to have their uniforms redone at least once per frame.
-		// Putting it here ensures that it's been after all render commands for the frame have prepared it.
-		((CCRenderStateMetal *)_renderState)->_uniformsPrepared = NO;
-	}
-	
 	CC_INCREMENT_GL_DRAWS(1);
 }
 

cocos2d/CCRenderer.m

@@ -712,11 +712,22 @@ -(void)dealloc
 
 @implementation CCGraphicsBufferBindings
 
-// Base implementations do do nothing.
--(void)prepare {}
--(void)commit {}
+// Base implementations does nothing.
 -(void)bind:(BOOL)bind {}
 
+-(void)prepare
+{
+	[_vertexBuffer prepare];
+	[_indexBuffer prepare];
+}
+
+-(void)commit
+{
+	[_vertexBuffer commit];
+	[_indexBuffer commit];
+}
+
+
 @end
 
 

cocos2d/CCRenderer_private.h

@@ -189,15 +189,15 @@ CCGraphicsBufferPushElements(CCGraphicsBuffer *buffer, size_t requestedCount)
 	CCGraphicsBuffer *_uniformBuffer;
 }
 
+/// Make the buffers ready to use by drawing commands.
+-(void)bind:(BOOL)bind;
+
 /// Prepare buffers for changes.
 -(void)prepare;
 
 /// Commit changes to buffers.
 -(void)commit;
 
-/// Bind the buffers. (Not used by Metal)
--(void)bind:(BOOL)bind;
-
 @end
 
 

cocos2d/CCShader.m

@@ -201,7 +201,7 @@ @implementation CCShader {
 
 //MARK: GL Uniform Setters:
 
-static CCGLUniformSetter
+static CCUniformSetter
 GLUniformSetFloat(NSString *name, GLint location)
 {
 	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
@@ -212,7 +212,7 @@ @implementation CCShader {
 	};
 }
 
-static CCGLUniformSetter
+static CCUniformSetter
 GLUniformSetVec2(NSString *name, GLint location)
 {
 	NSString *textureName = nil;
@@ -255,7 +255,7 @@ @implementation CCShader {
 	};
 }
 
-static CCGLUniformSetter
+static CCUniformSetter
 GLUniformSetVec3(NSString *name, GLint location)
 {
 	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
@@ -268,7 +268,7 @@ @implementation CCShader {
 	};
 }
 
-static CCGLUniformSetter
+static CCUniformSetter
 GLUniformSetVec4(NSString *name, GLint location)
 {
 	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
@@ -288,7 +288,7 @@ @implementation CCShader {
 	};
 }
 
-static CCGLUniformSetter
+static CCUniformSetter
 GLUniformSetMat4(NSString *name, GLint location)
 {
 	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
@@ -370,14 +370,149 @@ -(instancetype)initWithGLProgram:(GLuint)program uniformSetters:(NSDictionary *)
 }
 
 #if __CC_METAL_SUPPORTED_AND_ENABLED
+
+static CCUniformSetter
+MetalUniformSetBuffer(NSString *name, MTLArgument *vertexArg, MTLArgument *fragmentArg)
+{
+	NSUInteger vertexIndex = vertexArg.index;
+	NSUInteger fragmentIndex = fragmentArg.index;
+	size_t bytes = vertexArg.bufferDataSize;
+	
+	CCMetalContext *context = [CCMetalContext currentContext];
+	
+	// Handle cc_VertexAttributes specially.
+	if([name isEqualToString:@"cc_VertexAttributes"]){
+		NSCAssert(vertexArg && !fragmentArg, @"cc_VertexAttributes should only be used by vertex functions.");
+		NSCAssert(bytes == sizeof(CCVertex), @"cc_VertexAttributes data size is not sizeof(CCVertex).");
+		
+		return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
+			CCGraphicsBufferMetal *vertexBuffer = (CCGraphicsBufferMetal *)renderer->_buffers->_vertexBuffer;
+			id<MTLBuffer> metalBuffer = vertexBuffer->_buffer;
+			
+			[context->_currentRenderCommandEncoder setVertexBuffer:metalBuffer offset:0 atIndex:vertexIndex];
+		};
+	} else {
+		// If both args are active, they must match.
+		NSCAssert(!vertexArg || !fragmentArg || bytes == fragmentArg.bufferDataSize, @"Vertex and fragment argument type don't match for '%@'.", vertexArg.name);
+		
+		return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
+			CCGraphicsBufferMetal *uniformBuffer = (CCGraphicsBufferMetal *)renderer->_buffers->_uniformBuffer;
+			id<MTLBuffer> metalBuffer = uniformBuffer->_buffer;
+			
+			NSNumber *globalOffset = renderer->_globalShaderUniformBufferOffsets[name];
+			NSUInteger offset = globalOffset.unsignedIntegerValue;
+			
+			if(!globalOffset){
+				void *buff = CCGraphicsBufferPushElements(uniformBuffer, bytes);
+				[shaderUniforms[name] getValue:buff];
+				
+				offset = buff - uniformBuffer->_ptr;
+			}
+			
+			id<MTLRenderCommandEncoder> renderEncoder = context->_currentRenderCommandEncoder;
+			if(vertexArg) [renderEncoder setVertexBuffer:metalBuffer offset:offset atIndex:vertexIndex];
+			if(fragmentArg) [renderEncoder setFragmentBuffer:metalBuffer offset:offset atIndex:fragmentIndex];
+		};
+	}
+}
+
+static CCUniformSetter
+MetalUniformSetSampler(NSString *name, MTLArgument *vertexArg, MTLArgument *fragmentArg)
+{
+	NSUInteger vertexIndex = vertexArg.index;
+	NSUInteger fragmentIndex = fragmentArg.index;
+	
+	// For now, samplers and textures are locked together like in GL.
+	NSString *textureName = [name substringToIndex:name.length - @"Sampler".length];
+	
+	CCMetalContext *context = [CCMetalContext currentContext];
+	
+	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
+		CCTexture *texture = shaderUniforms[name] ?: globalShaderUniforms[name] ?: [CCTexture none];
+		NSCAssert([texture isKindOfClass:[CCTexture class]], @"Shader uniform '%@' value must be a CCTexture object.", name);
+		
+		id<MTLSamplerState> sampler = texture.metalSampler;
+		
+		id<MTLRenderCommandEncoder> renderEncoder = context->_currentRenderCommandEncoder;
+		if(vertexArg) [renderEncoder setVertexSamplerState:sampler atIndex:vertexIndex];
+		if(fragmentArg) [renderEncoder setFragmentSamplerState:sampler atIndex:fragmentIndex];
+	};
+}
+
+static CCUniformSetter
+MetalUniformSetTexture(NSString *name, MTLArgument *vertexArg, MTLArgument *fragmentArg)
+{
+	NSUInteger vertexIndex = vertexArg.index;
+	NSUInteger fragmentIndex = fragmentArg.index;
+	
+	CCMetalContext *context = [CCMetalContext currentContext];
+	
+	return ^(CCRenderer *renderer, NSDictionary *shaderUniforms, NSDictionary *globalShaderUniforms){
+		CCTexture *texture = shaderUniforms[name] ?: globalShaderUniforms[name] ?: [CCTexture none];
+		NSCAssert([texture isKindOfClass:[CCTexture class]], @"Shader uniform '%@' value must be a CCTexture object.", name);
+		
+		id<MTLTexture> metalTexture = texture.metalTexture;
+		
+		id<MTLRenderCommandEncoder> renderEncoder = context->_currentRenderCommandEncoder;
+		if(vertexArg) [renderEncoder setVertexTexture:metalTexture atIndex:vertexIndex];
+		if(fragmentArg) [renderEncoder setFragmentTexture:metalTexture atIndex:fragmentIndex];
+	};
+}
+
+static NSDictionary *
+MetalUniformSettersForFunctions(id<MTLFunction> vertexFunction, id<MTLFunction> fragmentFunction)
+{
+	// Get the shader reflection information by making a dummy render pipeline state.
+	MTLRenderPipelineDescriptor *descriptor = [MTLRenderPipelineDescriptor new];
+	descriptor.vertexFunction = vertexFunction;
+	descriptor.fragmentFunction = fragmentFunction;
+	
+	NSError *error = nil;
+	MTLRenderPipelineReflection *reflection = nil;
+	[[CCMetalContext currentContext].device newRenderPipelineStateWithDescriptor:descriptor options:MTLPipelineOptionArgumentInfo reflection:&reflection error:&error];
+	
+	NSCAssert(!error, @"Error getting Metal shader arguments.");
+	
+	// Collect all of the arguments.
+	NSMutableDictionary *vertexArgs = [NSMutableDictionary dictionary];
+	for(MTLArgument *arg in reflection.vertexArguments){ if(arg.active){ vertexArgs[arg.name] = arg; }}
+	
+	NSMutableDictionary *fragmentArgs = [NSMutableDictionary dictionary];
+	for(MTLArgument *arg in reflection.fragmentArguments){ if(arg.active){ fragmentArgs[arg.name] = arg; }}
+	
+	NSSet *argSet = [[NSSet setWithArray:vertexArgs.allKeys] setByAddingObjectsFromArray:fragmentArgs.allKeys];
+	
+	// Make uniform setters.
+	NSMutableDictionary *uniformSetters = [NSMutableDictionary dictionary];
+	
+	for(NSString *name in argSet){
+		MTLArgument *vertexArg = vertexArgs[name];
+		MTLArgument *fragmentArg = fragmentArgs[name];
+		
+		// If neither argument is active. Skip.
+		if(!vertexArg.active && !fragmentArg.active) continue;
+		
+		MTLArgumentType type = (vertexArg ? vertexArg.type : fragmentArg.type);
+		NSCAssert(!vertexArg || !fragmentArg || type == fragmentArg.type, @"Vertex and fragment argument type don't match for '%@'.", name);
+		
+		switch(type){
+			case MTLArgumentTypeBuffer: uniformSetters[name] = MetalUniformSetBuffer(name, vertexArg, fragmentArg); break;
+			case MTLArgumentTypeSampler: uniformSetters[name] = MetalUniformSetSampler(name, vertexArg, fragmentArg); break;
+			case MTLArgumentTypeTexture: uniformSetters[name] = MetalUniformSetTexture(name, vertexArg, fragmentArg); break;
+			case MTLArgumentTypeThreadgroupMemory: NSCAssert(NO, @"Compute memory not supported. (yet?)"); break;
+		}
+	}
+	
+	return uniformSetters;
+}
+
 -(instancetype)initWithMetalVertexFunction:(id<MTLFunction>)vertexFunction fragmentFunction:(id<MTLFunction>)fragmentFunction
 {
 	if((self = [super init])){
 		_vertexFunction = vertexFunction;
 		_fragmentFunction = fragmentFunction;
 
-		#warning TODO setup _uniformSetters
-//		_uniformSetters = uniformSetters;
+		_uniformSetters = MetalUniformSettersForFunctions(vertexFunction, fragmentFunction);
 	}
 
 	return self;
@@ -458,6 +593,8 @@ +(void)initialize
 #if __CC_METAL_SUPPORTED_AND_ENABLED
 	if([CCConfiguration sharedConfiguration].graphicsAPI == CCGraphicsAPIMetal){
 		id<MTLLibrary> library = [CCMetalContext currentContext].library;
+		NSAssert(library, @"Metal shader library not found.");
+		
 		id<MTLFunction> vertex = [library newFunctionWithName:@"CCVertexFunctionDefault"];
 
 		CC_SHADER_POS_COLOR = [[self alloc] initWithMetalVertexFunction:vertex fragmentFunction:[library newFunctionWithName:@"CCFragmentFunctionDefaultColor"]];

cocos2d/CCShader_Private.h

@@ -1,19 +1,22 @@
-#import "ccMacros.h"
+#import "CCShader.h"
+
 
 #if __CC_METAL_SUPPORTED_AND_ENABLED
 #import <Metal/Metal.h>
+#import "CCMetalSupport_Private.h"
 #endif
 
-#import "CCShader.h"
 
 @class CCRenderer;
 
-typedef void (^CCGLUniformSetter)(
+
+typedef void (^CCUniformSetter)(
 	__unsafe_unretained CCRenderer *renderer,
 	__unsafe_unretained NSDictionary *shaderUniforms,
 	__unsafe_unretained NSDictionary *globalShaderUniforms
 );
 
+
 @interface CCShader() {
 	@public
 	GLuint _program;