metal_context.hpp

Go to the documentation of this file.

//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
 
#ifndef metal_context_h
#define metal_context_h
 
#import <Metal/Metal.h>
 
#include "random.hpp"
#include "timing.hpp"
#include "piecewise.hpp"
 
namespace gpu {
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
    template<bool SAFE_MATH=false>
    class metal_context {
    private:
        id<MTLDevice> device;
        id<MTLCommandQueue> queue;
        std::unordered_map<graph::leaf_node<float, SAFE_MATH> *, id<MTLBuffer>> kernel_arguments;
        std::unordered_map<void *, id<MTLTexture>> texture_arguments;
        id<MTLCommandBuffer> command_buffer;
        id<MTLLibrary> library;
        std::unordered_map<std::string, std::vector<MTLMutability>> bufferMutability;
 
    public:
        constexpr static size_t random_state_scale = 3000;
        constexpr static size_t random_state_size = 1024*random_state_scale;
 
        int remaining_const_memory;
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static size_t max_concurrency() {
            return MTLCopyAllDevices().count;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        static std::string device_type() {
            return "Metal GPU";
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        metal_context(const size_t index) :
        device([MTLCopyAllDevices() objectAtIndex:index]),
        queue([device newCommandQueue]) {}
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void compile(const std::string kernel_source,
                     std::vector<std::string> names,
                     const bool add_reduction=false) {
            NSError *error;
            library = [device newLibraryWithSource:[NSString stringWithCString:kernel_source.c_str()
                                                                      encoding:NSUTF8StringEncoding]
                                           options:compile_options()
                                             error:&error];
 
            if (error) {
                NSLog(@"%@", error);
            }
 
            if (jit::verbose) {
                std::cout << "Metal GPU info." << std::endl
                          << "  Max thread group memory : " << device.maxThreadgroupMemoryLength << std::endl
                          << "  Max thread per group    : " << device.maxThreadsPerThreadgroup.width << std::endl
                          << "  Device name             : " << [device.name cStringUsingEncoding:NSString.defaultCStringEncoding] << std::endl
                          << "  Architecture            : " << [device.architecture.name cStringUsingEncoding:NSString.defaultCStringEncoding] << std::endl
                          << "  Max buffer length       : " << device.maxBufferLength << std::endl
                          << "  Max working set         : " << device.recommendedMaxWorkingSetSize << std::endl;
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_kernel_call(const std::string kernel_name,
                                                     graph::input_nodes<float, SAFE_MATH> inputs,
                                                     graph::output_nodes<float, SAFE_MATH> outputs,
                                                     graph::input_nodes<float, SAFE_MATH> atomics,
                                                     graph::shared_random_state<float, SAFE_MATH> state,
                                                     const size_t num_rays,
                                                     const jit::texture1d_list &tex1d_list,
                                                     const jit::texture2d_list &tex2d_list) {
            NSError *error;
 
            MTLFunctionDescriptor *funcDesc = [MTLFunctionDescriptor new];
            funcDesc.options = MTLFunctionOptionNone;
            funcDesc.name = [NSString stringWithCString:kernel_name.c_str()
                                               encoding:NSUTF8StringEncoding];
 
            id<MTLFunction> function = [library newFunctionWithDescriptor:funcDesc
                                                                    error:&error];
 
            MTLComputePipelineDescriptor *compute = [MTLComputePipelineDescriptor new];
            compute.threadGroupSizeIsMultipleOfThreadExecutionWidth = YES;
            compute.computeFunction = function;
            compute.maxTotalThreadsPerThreadgroup = 1024;
            for (size_t i = 0, ie = bufferMutability[kernel_name].size(); i < ie; i++) {
                compute.buffers[i].mutability = bufferMutability[kernel_name][i];
            }
 
            id<MTLComputePipelineState> pipeline = [device newComputePipelineStateWithDescriptor:compute
                                                                                         options:MTLPipelineOptionNone
                                                                                      reflection:NULL
                                                                                           error:&error];
 
            if (error) {
                NSLog(@"%@", error);
            }
 
            std::vector<id<MTLBuffer>> buffers;
            std::unordered_set<graph::leaf_node<float, SAFE_MATH> *> needed_buffers;
 
            const size_t buffer_element_size = sizeof(float);
            for (graph::shared_variable<float, SAFE_MATH> &input : inputs) {
                if (!kernel_arguments.contains(input.get())) {
                    backend::buffer<float> buffer = input->evaluate();
                    kernel_arguments[input.get()] = [device newBufferWithBytes:buffer.data()
                                                                        length:buffer.size()*buffer_element_size
                                                                       options:MTLResourceStorageModeShared];
                }
                if (!needed_buffers.contains(input.get())) {
                    buffers.push_back(kernel_arguments[input.get()]);
                    needed_buffers.insert(input.get());
                }
            }
            for (graph::shared_leaf<float, SAFE_MATH> &output : outputs) {
                if (!graph::atomic_accumulate_1D_cast(output).get()) {
                    if (!kernel_arguments.contains(output.get())) {
                        kernel_arguments[output.get()] = [device newBufferWithLength:num_rays*sizeof(float)
                                                                             options:MTLResourceStorageModeShared];
                    }
                    if (!needed_buffers.contains(output.get())) {
                        buffers.push_back(kernel_arguments[output.get()]);
                        needed_buffers.insert(output.get());
                    }
                }
            }
            for (graph::shared_variable<float, SAFE_MATH> &atomic : atomics) {
                if (!kernel_arguments.contains(atomic.get())) {
                    kernel_arguments[atomic.get()] = [device newBufferWithLength:atomic->size()*buffer_element_size
                                                                         options:MTLResourceStorageModeShared];
                }
                if (!needed_buffers.contains(atomic.get())) {
                    buffers.push_back(kernel_arguments[atomic.get()]);
                    needed_buffers.insert(atomic.get());
                }
            }
            if (state.get()) {
                if (!kernel_arguments.contains(state.get())) {
                    kernel_arguments[state.get()] = [device newBufferWithBytes:state->data()
                                                                        length:state->get_size_bytes()
                                                                       options:MTLResourceCPUCacheModeWriteCombined |
                                                                               MTLResourceStorageModeShared         |
                                                                               MTLResourceHazardTrackingModeUntracked];
                }
                buffers.push_back(kernel_arguments[state.get()]);
            }
 
            std::vector<id<MTLTexture>> textures;
            command_buffer = [queue commandBuffer];
            id<MTLBlitCommandEncoder> encoder = [command_buffer blitCommandEncoder];
            for (auto &[data, size] : tex1d_list) {
                if (!texture_arguments.contains(data)) {
                    MTLTextureDescriptor *descriptor = [MTLTextureDescriptor new];
                    descriptor.textureType = MTLTextureType1D;
                    descriptor.pixelFormat = MTLPixelFormatR32Float;
                    descriptor.width = size;
                    descriptor.storageMode = MTLStorageModeShared;
                    descriptor.cpuCacheMode = MTLCPUCacheModeWriteCombined;
                    descriptor.hazardTrackingMode = MTLHazardTrackingModeUntracked;
                    descriptor.usage = MTLTextureUsageShaderRead;
                    texture_arguments[data] = [device newTextureWithDescriptor:descriptor];
                    [texture_arguments[data] replaceRegion:MTLRegionMake1D(0, size)
                                               mipmapLevel:0
                                                 withBytes:reinterpret_cast<float *> (data)
                                               bytesPerRow:4*size];
 
                    [encoder optimizeContentsForGPUAccess:texture_arguments[data]];
                }
                textures.push_back(texture_arguments[data]);
            }
            for (auto &[data, size] : tex2d_list) {
                if (!texture_arguments.contains(data)) {
                    MTLTextureDescriptor *descriptor = [MTLTextureDescriptor new];
                    descriptor.textureType = MTLTextureType2D;
                    descriptor.pixelFormat = MTLPixelFormatR32Float;
                    descriptor.width = size[1];
                    descriptor.height = size[0];
                    descriptor.storageMode = MTLStorageModeShared;
                    descriptor.cpuCacheMode = MTLCPUCacheModeWriteCombined;
                    descriptor.hazardTrackingMode = MTLHazardTrackingModeUntracked;
                    descriptor.usage = MTLTextureUsageShaderRead;
                    texture_arguments[data] = [device newTextureWithDescriptor:descriptor];
                    [texture_arguments[data] replaceRegion:MTLRegionMake2D(0, 0, size[1], size[0])
                                               mipmapLevel:0
                                                 withBytes:reinterpret_cast<float *> (data)
                                               bytesPerRow:4*size[1]];
 
                    [encoder optimizeContentsForGPUAccess:texture_arguments[data]];
                }
                textures.push_back(texture_arguments[data]);
            }
            [encoder endEncoding];
            [command_buffer commit];
 
            std::vector<NSUInteger> offsets(buffers.size(), 0);
            NSRange range = NSMakeRange(0, buffers.size());
            NSRange tex_range = NSMakeRange(0, textures.size());
 
            NSUInteger total_parallel = state.get() ? state->size() : num_rays;
            NSUInteger thread_width = pipeline.threadExecutionWidth;
            NSUInteger threads_per_group = total_parallel < pipeline.maxTotalThreadsPerThreadgroup ?
                                                            thread_width                           :
                                                            pipeline.maxTotalThreadsPerThreadgroup;
            NSUInteger thread_groups = total_parallel/threads_per_group + (total_parallel%threads_per_group ? 1 : 0);
 
            if (jit::verbose) {
                std::cout << "  Kernel name : " << kernel_name << std::endl
                          << "    Thread execution width  : " << thread_width << std::endl
                          << "    Threads per group       : " << threads_per_group << std::endl
                          << "    Number of groups        : " << thread_groups << std::endl
                          << "    Total problem size      : " << threads_per_group*thread_groups << std::endl
                          << "    Total parallel size     : " << total_parallel << std::endl
                          << "    Current allocation size : " << device.currentAllocatedSize << std::endl;
            }
 
            if (state.get()) {
 
                return [this, num_rays, pipeline, buffers, offsets, range, tex_range, thread_groups, threads_per_group, textures
#ifdef PROFILE_KERNELS
                        , kernel_name
#endif
                ] () mutable {
                    command_buffer = [queue commandBuffer];
                    for (NSUInteger i = 0, ie = thread_groups*threads_per_group; i < num_rays; i += ie) {
                        id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoderWithDispatchType:MTLDispatchTypeSerial];
 
                        [encoder setComputePipelineState:pipeline];
                        [encoder setBuffers:buffers.data()
                                    offsets:offsets.data()
                                  withRange:range];
                        [encoder setBytes:&i
                                   length:sizeof(uint32_t)
                                  atIndex:buffers.size()];
                        [encoder setTextures:textures.data()
                                   withRange:tex_range];
 
                        [encoder dispatchThreadgroups:MTLSizeMake(thread_groups, 1, 1)
                                threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)];
                        [encoder endEncoding];
                    }
#ifdef PROFILE_KERNELS
                    [command_buffer addCompletedHandler:[kernel_name](id<MTLCommandBuffer> commandBuffer) {
                        std::cout << std::endl << "  " << kernel_name << " : " << commandBuffer.GPUEndTime - commandBuffer.GPUStartTime << " s" << std::endl << std::endl;
                    }];
#endif
                    [command_buffer commit];
                };
            } else {
                return [this, pipeline, buffers, offsets, range, tex_range, thread_groups, threads_per_group, textures
#ifdef PROFILE_KERNELS
                        , kernel_name
#endif
                ] () mutable {
                    command_buffer = [queue commandBuffer];
                    id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoderWithDispatchType:MTLDispatchTypeSerial];
 
                    [encoder setComputePipelineState:pipeline];
                    [encoder setBuffers:buffers.data()
                                offsets:offsets.data()
                              withRange:range];
                    [encoder setTextures:textures.data()
                               withRange:tex_range];
 
                    [encoder dispatchThreadgroups:MTLSizeMake(thread_groups, 1, 1)
                            threadsPerThreadgroup:MTLSizeMake(threads_per_group, 1, 1)];
                    [encoder endEncoding];
#ifdef PROFILE_KERNELS
                    [command_buffer addCompletedHandler:[kernel_name](id<MTLCommandBuffer> commandBuffer) {
                         std::cout << std::endl << "  " << kernel_name << " : " << commandBuffer.GPUEndTime - commandBuffer.GPUStartTime << " s" << std::endl << std::endl;
                    }];
#endif
                    [command_buffer commit];
                };
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<float(void)> create_max_call(graph::shared_leaf<float, SAFE_MATH> &argument,
                                                   std::function<void(void)> run) {
            NSError *error;
 
            MTLFunctionDescriptor *funcDesc = [MTLFunctionDescriptor new];
            funcDesc.options = MTLFunctionOptionNone;
            funcDesc.name = @"max_reduction";
 
            id<MTLFunction> function = [library newFunctionWithDescriptor:funcDesc
                                                                    error:&error];
 
            MTLComputePipelineDescriptor *compute = [MTLComputePipelineDescriptor new];
            compute.threadGroupSizeIsMultipleOfThreadExecutionWidth = YES;
            compute.computeFunction = function;
            compute.maxTotalThreadsPerThreadgroup = 1024;
            compute.buffers[0].mutability = MTLMutabilityImmutable;
 
            id<MTLComputePipelineState> max_state = [device newComputePipelineStateWithDescriptor:compute
                                                                                          options:MTLPipelineOptionNone
                                                                                       reflection:NULL
                                                                                            error:&error];
            if (error) {
                NSLog(@"%@", error);
            }
 
            id<MTLBuffer> result = [device newBufferWithLength:sizeof(float)
                                                       options:MTLResourceStorageModeShared];
 
            id<MTLBuffer> buffer = kernel_arguments[argument.get()];
 
            NSUInteger threads_per_group = max_state.maxTotalThreadsPerThreadgroup;
            NSUInteger thread_width = max_state.threadExecutionWidth;
            if (jit::verbose) {
                std::cout << "  Kernel name : max_reduction" << std::endl;
                std::cout << "    Thread execution width : " << thread_width << std::endl;
                std::cout << "    Threads per group      : " << threads_per_group << std::endl;
                std::cout << "    Number of groups       : " << 1 << std::endl;
                std::cout << "    Total problem size     : " << threads_per_group*1 << std::endl;
            }
 
            return [this, run, buffer, result, max_state] () mutable {
                run();
                command_buffer = [queue commandBuffer];
 
                id<MTLComputeCommandEncoder> encoder = [command_buffer computeCommandEncoderWithDispatchType:MTLDispatchTypeSerial];
 
                [encoder setComputePipelineState:max_state];
                [encoder setBuffer:buffer offset:0 atIndex:0];
                [encoder setBuffer:result offset:0 atIndex:1];
                [encoder dispatchThreadgroups:MTLSizeMake(1, 1, 1)
                        threadsPerThreadgroup:MTLSizeMake(1024, 1, 1)];
                [encoder endEncoding];
 
                [command_buffer commit];
                [command_buffer waitUntilCompleted];
 
                return static_cast<float *> (result.contents)[0];
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_zero_call(graph::input_nodes<float, SAFE_MATH> &inputs) {
            std::vector<id<MTLBuffer>> buffers;
            for (auto &input : inputs) {
                if (!kernel_arguments.contains(input.get())) {
                    kernel_arguments[input.get()] = [device newBufferWithBytes:input->data()
                                                                        length:input->size()*sizeof(float)
                                                                       options:MTLResourceStorageModeShared];
                }
                buffers.push_back(kernel_arguments[input.get()]);
            }
 
            std::vector<NSRange> ranges;
            for (id<MTLBuffer> buffer : buffers) {
                ranges.push_back(NSMakeRange(0, buffer.length));
            }
 
            return [this, buffers, ranges] () mutable {
                command_buffer = [queue commandBuffer];
                id<MTLBlitCommandEncoder> encoder = [command_buffer blitCommandEncoder];
 
                for (size_t i = 0, ie = buffers.size(); i < ie; i++) {
                    [encoder fillBuffer:buffers[i]
                                  range:ranges[i]
                                  value:0];
                }
                [encoder endEncoding];
#ifdef PROFILE_KERNELS
                [command_buffer addCompletedHandler:[](id<MTLCommandBuffer> commandBuffer) {
                    std::cout << std::endl << "  zero buffer : " << commandBuffer.GPUEndTime - commandBuffer.GPUStartTime << " s" << std::endl << std::endl;
                }];
#endif
                [command_buffer commit];
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> create_copy_call(graph::copy_nodes<float, SAFE_MATH> &setters) {
            std::vector<id<MTLBuffer>> sources;
            std::vector<id<MTLBuffer>> destinations;
 
            for (auto &[out, in] : setters) {
                if (!kernel_arguments.contains(in.get())) {
                    kernel_arguments[in.get()] = [device newBufferWithBytes:in->data()
                                                                     length:in->size()*sizeof(float)
                                                                    options:MTLResourceStorageModeShared];
                }
                destinations.push_back(kernel_arguments[in.get()]);
 
                if (!kernel_arguments.contains(out.get())) {
                    kernel_arguments[out.get()] = [device newBufferWithBytes:out->data()
                                                                      length:out->size()*sizeof(float)
                                                                     options:MTLResourceStorageModeShared];
                }
                sources.push_back(kernel_arguments[out.get()]);
            }
 
            return [this, sources, destinations] () mutable {
                command_buffer = [queue commandBuffer];
                id<MTLBlitCommandEncoder> encoder = [command_buffer blitCommandEncoder];
 
                for (size_t i = 0, ie = sources.size(); i < ie; i++) {
                    [encoder copyFromBuffer:sources[i]
                               sourceOffset:0
                                   toBuffer:destinations[i]
                          destinationOffset:0
                                       size:sources[i].length];
                }
                [encoder endEncoding];
#ifdef PROFILE_KERNELS
                [command_buffer addCompletedHandler:[](id<MTLCommandBuffer> commandBuffer) {
                    std::cout << std::endl << "  copy buffer : " << commandBuffer.GPUEndTime - commandBuffer.GPUStartTime << " s" << std::endl << std::endl;
                }];
#endif
                [command_buffer commit];
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        MTLCompileOptions *compile_options() {
            MTLCompileOptions *options = [MTLCompileOptions new];
            options.mathMode = MTLMathModeFast;
            options.mathFloatingPointFunctions = MTLMathFloatingPointFunctionsFast;
            options.optimizationLevel = MTLLibraryOptimizationLevelDefault;
            options.languageVersion = MTLLanguageVersion3_2;
            return options;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void wait() {
            command_buffer = [queue commandBuffer];
 
            [command_buffer commit];
            [command_buffer waitUntilCompleted];
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        std::function<void(void)> run_function(std::function<void(void)> callback) {
            return [this, callback]() {
                command_buffer = [queue commandBuffer];
        
                [command_buffer addCompletedHandler:[callback](id<MTLCommandBuffer> commandBuffer) {
#ifdef PROFILE_KERNELS
                    timing::measure_diagnostic timer("callback");
#endif
                    callback();
#ifdef PROFILE_KERNELS
                    timer.print();
#endif
                }];
        
                [command_buffer commit];
            };
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void print_results(const size_t index,
                           const graph::output_nodes<float, SAFE_MATH> &nodes) {
            wait();
            for (auto &out : nodes) {
                std::cout << static_cast<float *> ([kernel_arguments[out.get()] contents])[index] << " ";
            }
            std::cout << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        float check_value(const size_t index,
                          const graph::shared_leaf<float, SAFE_MATH> &node) {
            wait();
            return static_cast<float *> ([kernel_arguments[node.get()] contents])[index];
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_device(graph::shared_leaf<float, SAFE_MATH> node,
                            float *source) {
            const size_t size = [kernel_arguments[node.get()] length];
            memcpy([kernel_arguments[node.get()] contents],
                   source, size);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void copy_to_host(graph::shared_leaf<float, SAFE_MATH> node,
                          float *destination) {
            command_buffer = [queue commandBuffer];
 
            [command_buffer commit];
            [command_buffer waitUntilCompleted];
 
            memcpy(destination,
                   kernel_arguments[node.get()].contents,
                   kernel_arguments[node.get()].length);
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_header(std::ostringstream &source_buffer) {
            source_buffer << "#include <metal_stdlib>" << std::endl;
            source_buffer << "#include <metal_simdgroup>" << std::endl;
            source_buffer << "using namespace metal;" << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_kernel_prefix(std::ostringstream &source_buffer,
                                  const std::string name,
                                  graph::input_nodes<float, SAFE_MATH> &inputs,
                                  graph::output_nodes<float, SAFE_MATH> &outputs,
                                  graph::input_nodes<float, SAFE_MATH> atomics,
                                  graph::shared_random_state<float, SAFE_MATH> &state,
                                  const size_t size,
                                  const std::vector<bool> &is_constant,
                                  jit::register_map &registers,
                                  const jit::register_usage &usage,
                                  jit::texture1d_list &textures1d,
                                  jit::texture2d_list &textures2d,
                                  jit::argument_set &thread_shared,
                                  jit::register_map &thread_mem,
                                  const size_t iterations=1) {
            source_buffer << std::endl;
            source_buffer << "kernel void " << name << "(" << std::endl;
 
            bufferMutability[name] = std::vector<MTLMutability> ();
 
            size_t used_thread_mem = 0;
 
            size_t buffer_count = 0;
            jit::argument_set used_args;
            for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                if (!used_args.contains(inputs[i].get())) {
#ifdef USE_INPUT_CACHE
                    if (!is_constant[i] && iterations > 1) {
                        const size_t needed_mem = inputs[i]->size() > 1024 ?
                                                  1024*4                   :
                                                  32*4;
                        if (used_thread_mem + needed_mem < device.maxThreadgroupMemoryLength) {
                            used_thread_mem += needed_mem;
                            thread_shared.insert(inputs[i].get());
                        }
                    } else if (is_constant[i]           &&
                               inputs[i]->size() < size &&
                               inputs[i]->size() < 1024) {
                        const size_t needed_mem = inputs[i]->size()*4;
                        if (used_thread_mem + needed_mem < device.maxThreadgroupMemoryLength) {
                            used_thread_mem += needed_mem;
                            thread_shared.insert(inputs[i].get());
                            thread_mem[inputs[i].get()] = jit::to_string('t', inputs[i].get());
                        }
                    }
#endif
                    bufferMutability[name].push_back(is_constant[i] ? MTLMutabilityMutable : MTLMutabilityImmutable);
                    source_buffer << "    " << (is_constant[i] ? "constant" : "device")
                                  << " float *"
                                  << jit::to_string('v', inputs[i].get())
                                  << " [[buffer(" << buffer_count++ << ")]]";
                    inputs[i]->endline(source_buffer, usage, ',');
                    used_args.insert(inputs[i].get());
                }
            }
            assert(used_args.size() == inputs.size() &&
                   "Kernel inputs contain duplicates.");
 
            for (size_t i = 0, ie = outputs.size(); i < ie; i++) {
                if (!used_args.contains(outputs[i].get()) &&
                    !graph::atomic_accumulate_1D_cast(outputs[i]).get()) {
                    bufferMutability[name].push_back(MTLMutabilityMutable);
                    source_buffer << "    device float *"
                                  << jit::to_string('o', outputs[i].get())
                                  << " [[buffer(" << buffer_count++ << ")]],"
                                  << std::endl;
                    used_args.insert(outputs[i].get());
                }
            }
            assert(used_args.size() == inputs.size() + outputs.size() &&
                   "Kernel outputs contain duplicates.");
 
            for (size_t i = 0, ie = atomics.size(); i < ie; i++) {
                if (!used_args.contains(atomics[i].get())) {
                    bufferMutability[name].push_back(MTLMutabilityMutable);
                    source_buffer << "    device atomic_float *"
                                  << jit::to_string('v', atomics[i].get())
                                  << " [[buffer(" << buffer_count++ << ")]],"
                                  << std::endl;
                    used_args.insert(atomics[i].get());
                }
            }
            assert(used_args.size() == inputs.size() + outputs.size() +
                                       atomics.size() &&
                   "Kernel atomics contain duplicates.");
 
            if (state.get()) {
                bufferMutability[name].push_back(MTLMutabilityMutable);
                source_buffer << "    device mt_state *"
                              << jit::to_string('s', state.get())
                              << " [[buffer(" << buffer_count++ << ")]],"
                              << std::endl
                              << "    constant uint32_t &offset [[buffer("
                              << buffer_count++ << ")]],"
                              << std::endl;
            }
            size_t index = 0;
            for (auto &[key, value] : textures1d) {
                source_buffer << "    const texture1d<float, access::read> "
                              << jit::to_string('a', key)
                              << " [[texture(" << index++ << ")]],"
                              << std::endl;
            }
            for (auto &[key, value] : textures2d) {
                source_buffer << "    const texture2d<float, access::read> "
                              << jit::to_string('a', key)
                              << " [[texture(" << index++ << ")]],"
                              << std::endl;
            }
            if (thread_shared.size()) {
                source_buffer << "    ushort t_index [[thread_position_in_threadgroup]]," << std::endl;
                source_buffer << "    ushort t_total [[threads_per_threadgroup]]," << std::endl;
            }
            source_buffer << "    "
                          << jit::smallest_uint_type<float> (size)
                          << " index [[thread_position_in_grid]]) {" << std::endl
                          << "    if (";
            if (state.get()) {
                source_buffer << "offset + ";
            }
            source_buffer << "index < "  << size << ") {" << std::endl;
 
            for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                if (is_constant[i]) {
#ifdef USE_INPUT_CACHE
                    if (usage.at(inputs[i].get()) && inputs[i]->size() == size) {
                        registers[inputs[i].get()] = jit::to_string('r', inputs[i].get());
                        source_buffer << "        const ";
                        jit::add_type<float> (source_buffer);
                        source_buffer << " " << registers[inputs[i].get()] << " = "
                                      << jit::to_string('v', inputs[i].get())
                                      << "[";
                        if (state.get()) {
                            source_buffer << "offset + ";
                        }
                        source_buffer << "index]";
                        inputs[i]->endline(source_buffer, usage);
                    }
#else
                    registers[inputs[i].get()] = jit::to_string('v', inputs[i].get())
                                               + "["
                                               + (state.get() ? "offset + " : "")
                                               + "index]";
#endif
                }
            }
 
            if (thread_shared.size()) {
                for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                    if (thread_shared.contains(inputs[i].get()) && is_constant[i]) {
                        source_buffer << "        threadgroup float "
                                      << jit::to_string('t', inputs[i].get())
                                      << "[" << inputs[i]->size() << "]";
                        inputs[i]->endline(source_buffer, usage);
                    }
                }
                for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                    if (thread_shared.contains(inputs[i].get()) && is_constant[i]) {
                        source_buffer << "    for(int j = t_index; j < "
                                      << inputs[i]->size()
                                      << "; j += t_total) {" << std::endl
                                      << "        "
                                      << jit::to_string('t', inputs[i].get())
                                      << "[j] = "
                                      << jit::to_string('v', inputs[i].get())
                                      << "[j]";
                        inputs[i]->endline(source_buffer, usage);
                        source_buffer << "    }" << std::endl;
                    }
                }
                for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                    if (thread_shared.contains(inputs[i].get()) && is_constant[i]) {
                        source_buffer << "    threadgroup_barrier(mem_flags::mem_threadgroup);"
                                      << std::endl;
                        break;
                    }
                }
                for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                    if (thread_shared.contains(inputs[i].get()) && is_constant[i]) {
                        thread_shared.erase(inputs[i].get());
                    }
                }
                for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                    if (thread_shared.contains(inputs[i].get()) && !is_constant[i]) {
                        source_buffer << "        threadgroup float "
                                      << jit::to_string('t', inputs[i].get())
                                      << "["
                                      << (inputs[i]->size() > 1024 ? 1024 : 32)
                                      << "]";
                        inputs[i]->endline(source_buffer, usage);
                        source_buffer << "        "
                                      << jit::to_string('t', inputs[i].get())
                                      << "[t_index] = "
                                      << jit::to_string('v', inputs[i].get())
                                      << "[";
                        if (state.get()) {
                            source_buffer << "offset + ";
                        }
                        source_buffer << "index]";
                        inputs[i]->endline(source_buffer, usage);
                    }
                }
            }
 
            if (iterations > 1) {
                source_buffer << "    for (size_t j = 0; j < " << iterations << "; j++) {" << std::endl;
            }
 
            for (size_t i = 0, ie = inputs.size(); i < ie; i++) {
                if (!is_constant[i]) {
#ifdef USE_INPUT_CACHE
                    if (usage.at(inputs[i].get())) {
                        registers[inputs[i].get()] = jit::to_string('r', inputs[i].get());
                        source_buffer << "        const ";
                        jit::add_type<float> (source_buffer);
                        source_buffer << " " << registers[inputs[i].get()] << " = ";
                        if (thread_shared.contains(inputs[i].get())) {
                            source_buffer << jit::to_string('t', inputs[i].get())
                                          << "[t_index]";
                        } else {
                            source_buffer << jit::to_string('v', inputs[i].get())
                                          << "[";
                            if (state.get()) {
                                source_buffer << "offset + ";
                            }
                            source_buffer << "index]";
                        }
                        inputs[i]->endline(source_buffer, usage);
                    }
#else
                    if (thread_shared.contains(inputs[i].get())) {
                        registers[inputs[i].get()] = jit::to_string('t', inputs[i].get()) + "[t_index]";
                    } else {
                        registers[inputs[i].get()] = jit::to_string('v', inputs[i].get()) + "["
                                                   + (state.get() ? "offset + " : "")
                                                   + "index]";
                    }
#endif
                }
            }
            if (state.get()) {
#ifdef USE_INPUT_CACHE
                registers[state.get()] = jit::to_string('r', state.get());
                source_buffer << "        device mt_state &" << registers[state.get()]
                              << " = " << jit::to_string('s', state.get())
                              << "[index]";
                state->endline(source_buffer, usage);
#else
                registers[state.get()] = jit::to_string('s', state.get())
                                       + "[index]";
#endif
            }
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_kernel_postfix(std::ostringstream &source_buffer,
                                   graph::output_nodes<float, SAFE_MATH> &outputs,
                                   graph::map_nodes<float, SAFE_MATH> &setters,
                                   graph::shared_random_state<float, SAFE_MATH> state,
                                   jit::register_map &registers,
                                   const jit::register_usage &usage,
                                   const jit::argument_set &thread_shared,
                                   jit::register_map &thread_mem,
                                   const size_t iterations=1) {
            jit::argument_set out_registers;
            for (auto &[out, in] : setters) {
                if (!out->is_match(in)) {
                    auto a = out->compile(source_buffer, registers,
                                          thread_mem, usage);
                    source_buffer << "        ";
                    if (thread_shared.contains(in.get())) {
                        source_buffer << jit::to_string('t',  in.get())
                                      << "[t_index] = ";
                    } else {
                        source_buffer << jit::to_string('v',  in.get())
                                      << "[";
                        if (state.get()) {
                            source_buffer << "offset + ";
                        }
                        source_buffer << "index] = ";
                    }
                    if constexpr (SAFE_MATH) {
                        source_buffer << "isnan(" << registers[a.get()]
                                      << ") ? 0.0 : ";
                    }
                    source_buffer << registers[a.get()] << ";" << std::endl;
                    out_registers.insert(out.get());
                }
            }
 
            for (auto &out : outputs) {
                if (!graph::variable_cast(out).get()             &&
                    !graph::atomic_accumulate_1D_cast(out).get() &&
                    !out_registers.contains(out.get())) {
                    auto a = out->compile(source_buffer, registers,
                                          thread_mem, usage);
                    source_buffer << "        " << jit::to_string('o',  out.get())
                                  << "[";
                    if (state.get()) {
                        source_buffer << "offset + ";
                    }
                    source_buffer << "index] = ";
                    if constexpr (SAFE_MATH) {
                        if (!graph::random_cast(a).get()) {
                            source_buffer << "isnan(" << registers[a.get()]
                                          << ") ? 0.0 : ";
                        }
                    }
                    source_buffer << registers[a.get()] << ";" << std::endl;
                    out_registers.insert(out.get());
                }
            }
 
            if (iterations > 1) {
                source_buffer << "    }" << std::endl;
            }
            for (auto &[out, in] : setters) {
                if (thread_shared.contains(in.get())) {
                    source_buffer << "        "
                                  << jit::to_string('v',  in.get())
                                  << "[";
                    if (state.get()) {
                        source_buffer << "offset + ";
                    }
                    source_buffer << "index] = "
                                  << jit::to_string('t',  in.get())
                                  << "[t_index];" << std::endl;
                }
            }
            source_buffer << "    }" << std::endl << "}" << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        void create_reduction(std::ostringstream &source_buffer,
                              const size_t size) {
            source_buffer << std::endl;
            source_buffer << "kernel void max_reduction(" << std::endl;
            source_buffer << "    constant float *input [[buffer(0)]]," << std::endl;
            source_buffer << "    device float *result [[buffer(1)]]," << std::endl;
            source_buffer << "    uint i [[thread_position_in_grid]]," << std::endl;
            source_buffer << "    uint j [[simdgroup_index_in_threadgroup]]," << std::endl;
            source_buffer << "    uint k [[thread_index_in_simdgroup]]) {" << std::endl;
            source_buffer << "    if (i < " << size << ") {" << std::endl;
            source_buffer << "        float sub_max = input[i];" << std::endl;
            source_buffer << "        for (size_t index = i + 1024; index < " << size <<"; index += 1024) {" << std::endl;
            source_buffer << "            sub_max = max(sub_max, input[index]);" << std::endl;
            source_buffer << "        }" << std::endl;
            source_buffer << "        threadgroup float thread_max[32];" << std::endl;
            source_buffer << "        thread_max[j] = simd_max(sub_max);" << std::endl;
            source_buffer << "        threadgroup_barrier(mem_flags::mem_threadgroup);" << std::endl;
            source_buffer << "        if (j == 0) {"  << std::endl;
            source_buffer << "            *result = simd_max(thread_max[k]);"  << std::endl;
            source_buffer << "        }"  << std::endl;
            source_buffer << "    }"  << std::endl;
            source_buffer << "}" << std::endl << std::endl;
        }
 
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
        float *get_buffer(graph::shared_leaf<float, SAFE_MATH> &node) {
            return static_cast<float *> ([kernel_arguments[node.get()] contents]);
        }
    };
}
 
#endif /* metal_context_h */