69 lines
2.1 KiB
C#
69 lines
2.1 KiB
C#
using System;
|
|
using Unity.Burst;
|
|
using Unity.Jobs;
|
|
using UnityEngine.Scripting;
|
|
|
|
[assembly: AlwaysLinkAssembly]
|
|
[assembly: BurstCompile(OptimizeFor = OptimizeFor.FastCompilation)]
|
|
|
|
namespace Unity.Barracuda
|
|
{
|
|
|
|
/// <summary>
|
|
/// Burst specific BLAS implementation
|
|
/// </summary>
|
|
[Preserve]
|
|
public class BurstBLAS : BLASPlugin
|
|
{
|
|
/// <inheritdoc/>
|
|
public bool IsNative()
|
|
{
|
|
return false; // not a native fast BLAS implementation
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public bool IsCurrentPlatformSupported()
|
|
{
|
|
try
|
|
{
|
|
// Sanity test if all the dependencies of the job are met at runtime
|
|
// Also prevent compiler from optimising this out
|
|
new BurstCPUOps.MatrixMultiplyJob();
|
|
}
|
|
catch (Exception e)
|
|
{
|
|
D.Log($"C# Job system not found. Disabling {this.GetType()}. Error: {e}");
|
|
return false;
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public unsafe void SGEMM(float* Ap, int AM, int AN, float* Bp, int BM, int BN, float* Cp, int CM, int CN,
|
|
int bs,
|
|
bool transposeA = false, bool transposeB = false)
|
|
{
|
|
var noDependencies = new JobHandle();
|
|
var fence = ScheduleSGEMM(noDependencies, Ap, AM, AN, Bp, BM, BN, Cp, CM, CN, bs, transposeA, transposeB);
|
|
fence.Complete();
|
|
}
|
|
|
|
/// <inheritdoc/>
|
|
public unsafe JobHandle ScheduleSGEMM(JobHandle dependsOn,
|
|
float* Ap, int AM, int AN, float* Bp, int BM, int BN, float* Cp, int CM, int CN,
|
|
int bs, // NOTE: bs (block size) is ignored
|
|
bool transposeA = false, bool transposeB = false)
|
|
{
|
|
var job = new BurstCPUOps.MatrixMultiplyJob();
|
|
job.A = Ap; job.AM = AM; job.AN = AN;
|
|
job.B = Bp; job.BM = BM; job.BN = BN;
|
|
job.C = Cp; job.CM = CM; job.CN = CN;
|
|
job.transposeA = transposeA;
|
|
job.transposeB = transposeB;
|
|
|
|
return job.Schedule(dependsOn);
|
|
}
|
|
}
|
|
}
|