From 41eb19e65a0974e23048bd7b3b1eb1e2f569b1d0 Mon Sep 17 00:00:00 2001 From: Yong He Date: Thu, 8 Dec 2022 14:56:20 -0800 Subject: Auto-diff for matrix operations. (#2559) Co-authored-by: Yong He --- source/slang/core.meta.slang | 6 +- source/slang/diff.meta.slang | 32 ++++++ source/slang/hlsl.meta.slang | 111 ++++++++++++++++++++- source/slang/slang-emit-c-like.cpp | 20 +++- tests/autodiff/matrix-arithmetic-fwd.slang | 41 ++++++++ .../matrix-arithmetic-fwd.slang.expected.txt | 5 + 6 files changed, 209 insertions(+), 6 deletions(-) create mode 100644 tests/autodiff/matrix-arithmetic-fwd.slang create mode 100644 tests/autodiff/matrix-arithmetic-fwd.slang.expected.txt diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang index ed80b3730..edb98b3c2 100644 --- a/source/slang/core.meta.slang +++ b/source/slang/core.meta.slang @@ -93,7 +93,11 @@ interface __BuiltinArithmeticType : __BuiltinType /// A type that can be used for logical/bitwise operations [sealed] [builtin] -interface __BuiltinLogicalType : __BuiltinType {} +interface __BuiltinLogicalType : __BuiltinType +{ + /// Initialize from a 32-bit signed integer value. + __init(int value); +} /// A type that logically has a sign (positive/negative/zero) [sealed] diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang index a97ab9eaf..248112810 100644 --- a/source/slang/diff.meta.slang +++ b/source/slang/diff.meta.slang @@ -110,6 +110,38 @@ void updatePair(inout DifferentialPair p, T newPrimal, T p = DifferentialPair(newPrimal, newDiff); } +// vector-matrix +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +DifferentialPair> mul(DifferentialPair> left, DifferentialPair> right) +{ + let primal = mul(left.p, right.p); + let diff = mul(left.d, right.p) + mul(left.p, right.d); + return DifferentialPair>(primal, diff); +} + +// matrix-vector +__generic +[ForceInline] +[ForwardDerivativeOf(mul)] +DifferentialPair> mul(DifferentialPair> left, DifferentialPair> right) +{ + let primal = mul(left.p, right.p); + let diff = mul(left.d, right.p) + mul(left.p, right.d); + return DifferentialPair>(primal, diff); +} + + +// matrix-matrix +__generic +[ForwardDerivativeOf(mul)] +DifferentialPair> mul(DifferentialPair> right, DifferentialPair> left) +{ + let primal = mul(right.p, left.p); + let diff = mul(right.d, left.p) + mul(right.p, left.d); + return DifferentialPair>(primal, diff); +} #define VECTOR_MAP_D_UNARY(TYPE, COUNT, D_FUNC, VALUE) \ vector result; \ diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang index 0232129ee..ed150a4c0 100644 --- a/source/slang/hlsl.meta.slang +++ b/source/slang/hlsl.meta.slang @@ -2715,7 +2715,24 @@ T mul(vector x, vector y) } // vector-matrix -__generic +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +vector mul(vector left, matrix right) +{ + vector result; + for( int j = 0; j < M; ++j ) + { + T sum = T(0); + for( int i = 0; i < N; ++i ) + { + sum += left[i] * right[i][j]; + } + result[j] = sum; + } + return result; +} +__generic __target_intrinsic(hlsl) __target_intrinsic(glsl, "($1 * $0)") vector mul(vector left, matrix right) @@ -2732,9 +2749,26 @@ vector mul(vector left, matrix right) } return result; } +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +vector mul(vector left, matrix right) +{ + vector result; + for( int j = 0; j < M; ++j ) + { + T sum = T(0); + for( int i = 0; i < N; ++i ) + { + sum |= left[i] & right[i][j]; + } + result[j] = sum; + } + return result; +} // matrix-vector -__generic +__generic __target_intrinsic(hlsl) __target_intrinsic(glsl, "($1 * $0)") vector mul(matrix left, vector right) @@ -2751,10 +2785,43 @@ vector mul(matrix left, vector right) } return result; } - +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +vector mul(matrix left, vector right) +{ + vector result; + for( int i = 0; i < N; ++i ) + { + T sum = T(0); + for( int j = 0; j < M; ++j ) + { + sum += left[i][j] * right[j]; + } + result[i] = sum; + } + return result; +} +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +vector mul(matrix left, vector right) +{ + vector result; + for( int i = 0; i < N; ++i ) + { + T sum = T(0); + for( int j = 0; j < M; ++j ) + { + sum |= left[i][j] & right[j]; + } + result[i] = sum; + } + return result; +} // matrix-matrix -__generic +__generic __target_intrinsic(hlsl) __target_intrinsic(glsl, "($1 * $0)") matrix mul(matrix right, matrix left) @@ -2772,6 +2839,42 @@ matrix mul(matrix right, matrix left) } return result; } +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +matrix mul(matrix right, matrix left) +{ + matrix result; + for( int r = 0; r < R; ++r) + for( int c = 0; c < C; ++c) + { + T sum = T(0); + for( int i = 0; i < N; ++i ) + { + sum += left[r][i] * right[i][c]; + } + result[r][c] = sum; + } + return result; +} +__generic +__target_intrinsic(hlsl) +__target_intrinsic(glsl, "($1 * $0)") +matrix mul(matrix right, matrix left) +{ + matrix result; + for( int r = 0; r < R; ++r) + for( int c = 0; c < C; ++c) + { + T sum = T(0); + for( int i = 0; i < N; ++i ) + { + sum |= left[r][i] & right[i][c]; + } + result[r][c] = sum; + } + return result; +} // noise (deprecated) diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp index cadab5690..dd40b7856 100644 --- a/source/slang/slang-emit-c-like.cpp +++ b/source/slang/slang-emit-c-like.cpp @@ -1778,7 +1778,6 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO case kIROp_MakeVector: case kIROp_MakeMatrix: - case kIROp_MakeMatrixFromScalar: case kIROp_MatrixReshape: case kIROp_VectorReshape: case kIROp_CastFloatToInt: @@ -1789,6 +1788,25 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO emitType(inst->getDataType()); emitArgs(inst); break; + case kIROp_MakeMatrixFromScalar: + { + emitType(inst->getDataType()); + auto matrixType = as(inst->getDataType()); + SLANG_RELEASE_ASSERT(matrixType); + auto columnCount = as(matrixType->getColumnCount()); + SLANG_RELEASE_ASSERT(columnCount); + auto rowCount = as(matrixType->getRowCount()); + SLANG_RELEASE_ASSERT(rowCount); + m_writer->emit("("); + for (IRIntegerValue i = 0; i < rowCount->getValue() * columnCount->getValue(); i++) + { + if (i != 0) + m_writer->emit(", "); + emitOperand(inst->getOperand(0), getInfo(EmitOp::General)); + } + m_writer->emit(")"); + } + break; case kIROp_AllocObj: m_writer->emit("new "); m_writer->emit(getName(inst->getDataType())); diff --git a/tests/autodiff/matrix-arithmetic-fwd.slang b/tests/autodiff/matrix-arithmetic-fwd.slang new file mode 100644 index 000000000..7a953cef8 --- /dev/null +++ b/tests/autodiff/matrix-arithmetic-fwd.slang @@ -0,0 +1,41 @@ +//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -shaderobj -output-using-type +//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj -output-using-type + +//TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name=outputBuffer +RWStructuredBuffer outputBuffer; + +[ForwardDifferentiable] +float3x3 g(float3x3 x, float3x3 y) +{ + float3x3 a = x + y; + float3x3 b = x - y; + return a * b + 2 * x * y; +} + +[ForwardDifferentiable] +float h(float2x2 x, float2x2 y) +{ + let t = mul(x, y); + return t[0][0] + t[0][1] + t[1][0] + t[1][1]; +} + +[numthreads(1, 1, 1)] +void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID) +{ + float3x3 a = float3x3(2.0); + float3x3 b = float3x3(1.5); + float3x3 da = float3x3(1.0); + + outputBuffer[0] = __fwd_diff(g)( + DifferentialPair(a, da), + DifferentialPair(b, da)).d._11; // Expect: 8 + + float2x2 l = float2x2(1.0, 2.0, 3.0, 4.0); + float2x2 r = float2x2(10.0, 11.0, 12.0, 13.0); + float2x2 d = float2x2(1.0, 0.0, 1.0, 1.0); + + //float2x2 epsilon = d * 0.001f; + //outputBuffer[1] = (h(l + epsilon, r + epsilon) - h(l - epsilon, r - epsilon)) / (epsilon[0][0] * 2.0)); + + outputBuffer[1] = __fwd_diff(h)(DifferentialPair(l, d), DifferentialPair(r, d)).d; // Expect 83.0 +} diff --git a/tests/autodiff/matrix-arithmetic-fwd.slang.expected.txt b/tests/autodiff/matrix-arithmetic-fwd.slang.expected.txt new file mode 100644 index 000000000..c595048c3 --- /dev/null +++ b/tests/autodiff/matrix-arithmetic-fwd.slang.expected.txt @@ -0,0 +1,5 @@ +type: float +8.0 +83.0 +0.0 +0.0 \ No newline at end of file -- cgit v1.2.3