Perkalian Matriks Berdasarkan Elemen Logam Apple (Produk Hadamard)

Baiklah, jawab pertanyaan saya sendiri di sini berdasarkan rekomendasi dari komentator - coba tulis shader saya sendiri!

Berikut kode shadernya:

#include <metal_stdlib>
using namespace metal;

/*
 hadamardProduct:
 Perform an element-wise multiplication (hadamard product) of the two input matrices A and B, store the result in C
 */
kernel void hadamardProductKernel(
texture_buffer<float, access::read> A [[texture(0)]],
texture_buffer<float, access::read> B [[texture(1)]],
texture_buffer<float, access::write> C [[texture(2)]],
uint gid [[thread_position_in_grid]]) {
    // C[i,j] = A[i,j] * B[i,j]
    C.write(A.read(gid) * B.read(gid), gid);
}

Dan swift yang mengeksekusi shader pada dua matriks 4x4:

import Foundation
import Metal
import MetalKit

guard
    let gpu = MTLCreateSystemDefaultDevice(),
    let commandQueue = gpu.makeCommandQueue(),
    let commandBuffer = commandQueue.makeCommandBuffer(),
    let defaultLibrary = gpu.makeDefaultLibrary(),
    let kernelFunction = defaultLibrary.makeFunction(name: "hadamardProductKernel")
else {exit(1)}

// Create the matrices to multiply (as row-major matrices)
var A:[Float] = [2,0,0,0,
                 0,2,0,0,
                 0,0,2,0,
                 0,0,0,2]

var B:[Float] = [1,0,0,0,
                 0,2,0,0,
                 0,0,3,0,
                 0,0,0,4]

let A_buffer = gpu.makeTexture(descriptor: MTLTextureDescriptor.textureBufferDescriptor(with: .r32Float,
                                                                                                                                                                                width: 16,
                                                                                                                                                                                resourceOptions: .storageModeManaged,
                                                                                                                                                                                usage: .shaderRead))
let B_buffer = gpu.makeTexture(descriptor: MTLTextureDescriptor.textureBufferDescriptor(with: .r32Float,
                                                                                                                                                                                width: 16,
                                                                                                                                                                                resourceOptions: .storageModeManaged,
                                                                                                                                                                                usage: .shaderRead))
let C_buffer = gpu.makeTexture(descriptor: MTLTextureDescriptor.textureBufferDescriptor(with: .r32Float,
                                                                                                                                                                                width: 16,
                                                                                                                                                                                resourceOptions: .storageModeManaged,
                                                                                                                                                                                usage: .shaderWrite))
A_buffer?.replace(region: MTLRegionMake1D(0, 16),
                  mipmapLevel: 0,
                  withBytes: UnsafeRawPointer(A),
                  bytesPerRow: 64)
B_buffer?.replace(region: MTLRegionMake1D(0, 16),
                  mipmapLevel: 0,
                  withBytes: UnsafeRawPointer(B),
                  bytesPerRow: 64)

let computePipelineState = try gpu.makeComputePipelineState(function: kernelFunction)
let computeEncoder = commandBuffer.makeComputeCommandEncoder()
computeEncoder?.setComputePipelineState(computePipelineState)
computeEncoder?.setTexture(A_buffer, index: 0)
computeEncoder?.setTexture(B_buffer, index: 1)
computeEncoder?.setTexture(C_buffer, index: 2)
let threadGroupSize = MTLSize(width: 16, height: 1, depth: 1)
let threadGroupCount = MTLSize(width: 1, height: 1, depth: 1)
computeEncoder?.dispatchThreadgroups(threadGroupCount, threadsPerThreadgroup: threadGroupSize)
computeEncoder?.endEncoding()
commandBuffer.commit()
commandBuffer.waitUntilCompleted()

print("done")

Hargai setiap komentar yang tertaut ke sumber daya untuk pembelajaran lebih lanjut tentang hal semacam ini.

Pilihan lainnya adalah menggunakan MTLBuffers (dalam contoh saya, saya menyimpan hasil di buffer input pertama):

#include <metal_stdlib>
using namespace metal;

kernel void hadamardProductKernel(
    device float *a [[ buffer(0) ]],
    const device float *b [[ buffer(1) ]],
    uint id [[ thread_position_in_grid ]]
)
{
    a[id] = a[id] * b[id];
}

Berikut kode Objective C yang mengeksekusi produk Hadamard pada dua array float32 (a->data dan b->data):

id<MTLLibrary> library = [device newDefaultLibrary];
id<MTLFunction> function = [library newFunctionWithName:@"hadamardProductKernel"];
id<MTLCommandBuffer> commandBuffer = [commandQueue commandBuffer];
id<MTLComputePipelineState> computePipelineState = [device newComputePipelineStateWithFunction:function error:NULL];
id<MTLComputeCommandEncoder> computeCommandEncoder = [commandBuffer computeCommandEncoder];
[computeCommandEncoder setComputePipelineState:computePipelineState];
id<MTLBuffer> buffer_a = (__bridge id<MTLBuffer>)(a->data);
[computeCommandEncoder setBuffer:buffer_a offset:0 atIndex:0];
id<MTLBuffer> buffer_b = (__bridge id<MTLBuffer>)(b->data);
[computeCommandEncoder setBuffer:buffer_b offset:0 atIndex:1];
MTLSize threadGroupSize = MTLSizeMake(<<ELEMENTS COUNT HERE>>, 1, 1);
MTLSize threadGroupCount = MTLSizeMake(1, 1, 1);
[computeCommandEncoder dispatchThreadgroups:threadGroupSize threadsPerThreadgroup:threadGroupCount];
[computeCommandEncoder endEncoding];
[commandBuffer commit];
[commandBuffer waitUntilCompleted];