From ddebd60853b3f34bfd8e89de804fd15808abf75d Mon Sep 17 00:00:00 2001 From: Yong He Date: Tue, 9 May 2023 18:00:48 -0700 Subject: [PATCH] Various fixes for autodiff and slangpy. (#2876) * Various fixes for autodiff and slangpy. * Fix cuda code gen for `select`. * Fix getBuildTagString(). * Fix. --------- Co-authored-by: Yong He --- prelude/slang-cpp-host-prelude.h | 16 +++++++ prelude/slang-cpp-types-core.h | 17 +++++++ prelude/slang-cuda-prelude.h | 32 +++++++++++++ prelude/slang-torch-prelude.h | 17 +++++++ source/slang/diff.meta.slang | 18 +++++++- source/slang/slang-diagnostic-defs.h | 3 ++ source/slang/slang-emit-c-like.cpp | 1 - source/slang/slang-emit-cpp.cpp | 11 +++++ .../slang-ir-check-differentiability.cpp | 46 +++++++++++++++++++ source/slang/slang-ir-util.cpp | 28 +++++------ source/slang/slang.cpp | 8 ++++ 11 files changed, 179 insertions(+), 18 deletions(-) diff --git a/prelude/slang-cpp-host-prelude.h b/prelude/slang-cpp-host-prelude.h index 0a026471c4..f69d03eed4 100644 --- a/prelude/slang-cpp-host-prelude.h +++ b/prelude/slang-cpp-host-prelude.h @@ -28,6 +28,22 @@ # include #endif // SLANG_LLVM +#if defined(_MSC_VER) +# define SLANG_PRELUDE_SHARED_LIB_EXPORT __declspec(dllexport) +#else +# define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__((__visibility__("default"))) +//# define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__ ((dllexport)) __attribute__((__visibility__("default"))) +#endif + +#ifdef __cplusplus +# define SLANG_PRELUDE_EXTERN_C extern "C" +# define SLANG_PRELUDE_EXTERN_C_START extern "C" { +# define SLANG_PRELUDE_EXTERN_C_END } +#else +# define SLANG_PRELUDE_EXTERN_C +# define SLANG_PRELUDE_EXTERN_C_START +# define SLANG_PRELUDE_EXTERN_C_END +#endif #include "slang-cpp-scalar-intrinsics.h" diff --git a/prelude/slang-cpp-types-core.h b/prelude/slang-cpp-types-core.h index c49ee013c4..25fe472021 100644 --- a/prelude/slang-cpp-types-core.h +++ b/prelude/slang-cpp-types-core.h @@ -202,6 +202,23 @@ struct Vector }; +template +SLANG_FORCE_INLINE Vector _slang_select(Vector condition, Vector v0, Vector v1) +{ + Vector result; + for (int i = 0; i < N; i++) + { + result[i] = condition[i] ? v0[i] : v1[i]; + } + return result; +} + +template +SLANG_FORCE_INLINE T _slang_select(bool condition, T v0, T v1) +{ + return condition ? v0 : v1; +} + template SLANG_FORCE_INLINE T _slang_vector_get_element(Vector x, int index) { diff --git a/prelude/slang-cuda-prelude.h b/prelude/slang-cuda-prelude.h index 0e0349bd79..0362a3a677 100644 --- a/prelude/slang-cuda-prelude.h +++ b/prelude/slang-cuda-prelude.h @@ -743,11 +743,43 @@ SLANG_FLOAT_MATRIX_OPS(__half) #undef SLANG_MATRIX_INT_NEG_OP #undef SLANG_FLOAT_MATRIX_MOD +#define SLANG_SELECT_IMPL(T, N)\ +SLANG_FORCE_INLINE SLANG_CUDA_CALL Vector _slang_select(bool##N condition, Vector v0, Vector v1) \ +{ \ + Vector result; \ + for (int i = 0; i < N; i++) \ + { \ + *_slang_vector_get_element_ptr(&result, i) = _slang_vector_get_element(condition, i) ? _slang_vector_get_element(v0, i) : _slang_vector_get_element(v1, i); \ + } \ + return result; \ +} +#define SLANG_SELECT_T(T)\ + SLANG_SELECT_IMPL(T, 2)\ + SLANG_SELECT_IMPL(T, 3)\ + SLANG_SELECT_IMPL(T, 4) + +SLANG_SELECT_T(int) +SLANG_SELECT_T(uint) +SLANG_SELECT_T(short) +SLANG_SELECT_T(ushort) +SLANG_SELECT_T(char) +SLANG_SELECT_T(uchar) +SLANG_SELECT_T(float) +SLANG_SELECT_T(double) + +template +SLANG_FORCE_INLINE SLANG_CUDA_CALL T _slang_select(bool condition, T v0, T v1) +{ + return condition ? v0 : v1; +} + // // Half support // #if SLANG_CUDA_ENABLE_HALF +SLANG_SELECT_T(__half) + // Convenience functions ushort -> half SLANG_FORCE_INLINE SLANG_CUDA_CALL __half2 __ushort_as_half(const ushort2& i) { return __halves2half2(__ushort_as_half(i.x), __ushort_as_half(i.y)); } diff --git a/prelude/slang-torch-prelude.h b/prelude/slang-torch-prelude.h index 2f5273e1fc..0a69b8305b 100644 --- a/prelude/slang-torch-prelude.h +++ b/prelude/slang-torch-prelude.h @@ -38,6 +38,23 @@ # include #endif // SLANG_LLVM +#if defined(_MSC_VER) +# define SLANG_PRELUDE_SHARED_LIB_EXPORT __declspec(dllexport) +#else +# define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__((__visibility__("default"))) +//# define SLANG_PRELUDE_SHARED_LIB_EXPORT __attribute__ ((dllexport)) __attribute__((__visibility__("default"))) +#endif + +#ifdef __cplusplus +# define SLANG_PRELUDE_EXTERN_C extern "C" +# define SLANG_PRELUDE_EXTERN_C_START extern "C" { +# define SLANG_PRELUDE_EXTERN_C_END } +#else +# define SLANG_PRELUDE_EXTERN_C +# define SLANG_PRELUDE_EXTERN_C_START +# define SLANG_PRELUDE_EXTERN_C_END +#endif + #define SLANG_PRELUDE_NAMESPACE diff --git a/source/slang/diff.meta.slang b/source/slang/diff.meta.slang index 0725103da5..f2fd8e3b0e 100644 --- a/source/slang/diff.meta.slang +++ b/source/slang/diff.meta.slang @@ -246,52 +246,66 @@ __intrinsic_type($(kIROp_TorchTensorType)) struct TorchTensor { __intrinsic_op($(kIROp_TorchTensorGetView)) + [CudaHost] TensorView getView(); __target_intrinsic(cuda, "$0.dims()") __target_intrinsic(cpp, "$0.dims()") [__readNone] + [CudaHost] uint dims(); __target_intrinsic(cuda, "$0.size($1)") __target_intrinsic(cpp, "$0.size($1)") [__readNone] + [CudaHost] uint size(uint i); __target_intrinsic(cuda, "$0.stride($1)") __target_intrinsic(cpp, "$0.stride($1)") [__readNone] + [CudaHost] uint stride(uint i); __target_intrinsic(cuda, "$0.data_ptr<$G0>()") __target_intrinsic(cpp, "$0.data_ptr<$G0>()") [__readNone] + [CudaHost] Ptr data_ptr(); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor alloc(uint x); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor alloc(uint x, uint y); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor alloc(uint x, uint y, uint z); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor alloc(uint x, uint y, uint z, uint w); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor alloc(uint i0, uint i1, uint i2, uint i3, uint i4); __intrinsic_op($(kIROp_AllocateTorchTensor)) + [CudaHost] static TorchTensor emptyLike(TorchTensor other); __target_intrinsic(cpp, "$0.zero_()") + [CudaHost] void fillZero(); __target_intrinsic(cpp, "$0.fill_($1)") + [CudaHost] void fillValue(T val); + [CudaHost] static TorchTensor zerosLike(TorchTensor other) { var result = emptyLike(other); @@ -854,8 +868,10 @@ T detach(T x); #define SLANG_SQR(x) ((x)*(x)) +#define SLANG_SIGN(x) select(((x)>T(0.0)), ReturnType(T(1.0)), select(((x)==T(0.0)), ReturnType(T(0.0)), ReturnType(T(-1.0)))) + // Absolute value -UNARY_DERIVATIVE_IMPL(abs, select(dpx.p > T(0.0), dpx.d, ReturnType.dmul(T(-1.0), dpx.d)), (ReturnType.dmul(__slang_noop_cast(sign(dpx.p)), dOut))) +UNARY_DERIVATIVE_IMPL(abs, ReturnType.dmul(SLANG_SIGN(dpx.p), dpx.d), ReturnType.dmul(SLANG_SIGN(dpx.p), dOut)) // Saturate UNARY_DERIVATIVE_IMPL(saturate, select(dpx.p < T(0.0) || dpx.p > T(1.0), ReturnType.dzero(), dpx.d), select(dpx.p < T(0.0) || dpx.p > T(1.0), ReturnType.dzero(), dOut)) // frac diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h index a11474c8f0..723dc4d64f 100644 --- a/source/slang/slang-diagnostic-defs.h +++ b/source/slang/slang-diagnostic-defs.h @@ -614,6 +614,9 @@ DIAGNOSTIC(41023, Error, getStringHashMustBeOnStringLiteral, "getStringHash can DIAGNOSTIC(41901, Error, unsupportedUseOfLValueForAutoDiff, "unsupported use of L-value for auto differentiation.") DIAGNOSTIC(41902, Error, cannotDifferentiateDynamicallyIndexedData, "cannot auto-differentiate mixed read/write access to dynamically indexed data in '$0'.") + +DIAGNOSTIC(42001, Error, invalidUseOfTorchTensorTypeInDeviceFunc, "invalid use of TorchTensor type in device/kernel functions. use `TensorView` instead.") + // // 5xxxx - Target code generation. // diff --git a/source/slang/slang-emit-c-like.cpp b/source/slang/slang-emit-c-like.cpp index 2b5323e3f0..75d7b3cf85 100644 --- a/source/slang/slang-emit-c-like.cpp +++ b/source/slang/slang-emit-c-like.cpp @@ -2215,7 +2215,6 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO case kIROp_Select: { - auto prec = getInfo(EmitOp::Conditional); needClose = maybeEmitParens(outerPrec, prec); diff --git a/source/slang/slang-emit-cpp.cpp b/source/slang/slang-emit-cpp.cpp index df6ab49013..49edec92b5 100644 --- a/source/slang/slang-emit-cpp.cpp +++ b/source/slang/slang-emit-cpp.cpp @@ -1615,6 +1615,17 @@ bool CPPSourceEmitter::tryEmitInstExprImpl(IRInst* inst, const EmitOpInfo& inOut m_writer->emit(")"); return true; } + case kIROp_Select: + { + m_writer->emit("_slang_select("); + emitOperand(inst->getOperand(0), getInfo(EmitOp::General)); + m_writer->emit(", "); + emitOperand(inst->getOperand(1), getInfo(EmitOp::General)); + m_writer->emit(","); + emitOperand(inst->getOperand(2), getInfo(EmitOp::General)); + m_writer->emit(")"); + return true; + } } } diff --git a/source/slang/slang-ir-check-differentiability.cpp b/source/slang/slang-ir-check-differentiability.cpp index d81a337195..6c9280fca5 100644 --- a/source/slang/slang-ir-check-differentiability.cpp +++ b/source/slang/slang-ir-check-differentiability.cpp @@ -195,8 +195,54 @@ struct CheckDifferentiabilityPassContext : public InstPassBase } } + bool checkType(IRInst* type) + { + type = unwrapAttributedType(type); + if (as(type)) + return false; + else if (auto arrayType = as(type)) + return checkType(arrayType->getElementType()); + else if (auto structType = as(type)) + { + for (auto field : structType->getFields()) + { + if (!checkType(field->getFieldType())) + return false; + } + } + return true; + } + void checkForInvalidHostTypeUsage(IRGlobalValueWithCode* funcInst) + { + auto outerFuncInst = maybeFindOuterGeneric(funcInst); + + if (outerFuncInst->findDecoration()) + return; + if (outerFuncInst->findDecoration()) + return; + + // This is a kernel function, we don't allow using TorchTensor type here. + for (auto b : funcInst->getBlocks()) + { + for (auto inst : b->getChildren()) + { + if (!checkType(inst->getDataType())) + { + auto loc = inst->sourceLoc; + if (!loc.isValid()) + loc = funcInst->sourceLoc; + sink->diagnose(loc, Diagnostics::invalidUseOfTorchTensorTypeInDeviceFunc); + return; + } + + } + } + } + void processFunc(IRGlobalValueWithCode* funcInst) { + checkForInvalidHostTypeUsage(funcInst); + if (!_isFuncMarkedForAutoDiff(funcInst)) return; if (!funcInst->getFirstBlock()) diff --git a/source/slang/slang-ir-util.cpp b/source/slang/slang-ir-util.cpp index ba34a725d7..55c2b18c0f 100644 --- a/source/slang/slang-ir-util.cpp +++ b/source/slang/slang-ir-util.cpp @@ -411,10 +411,6 @@ bool isPtrLikeOrHandleType(IRInst* type) return true; if (as(type)) return true; - if (as(type)) - return true; - if (as(type)) - return true; switch (type->getOp()) { case kIROp_ComPtrType: @@ -824,30 +820,30 @@ bool isGlobalOrUnknownMutableAddress(IRGlobalValueWithCode* parentFunc, IRInst* if (!isPtrLikeOrHandleType(type)) return false; + if (root) + { + if (as(root->getDataType())) + { + return false; + } + } + switch (root->getOp()) { case kIROp_GlobalVar: - return true; case kIROp_GlobalParam: case kIROp_GlobalConstant: case kIROp_Var: case kIROp_Param: break; + case kIROp_Call: + return true; default: - // The inst is defined by an unknown inst. return true; } - if (root) - { - if (as(root->getDataType())) - { - return false; - } - auto addrInstParent = getParentFunc(root); - return (addrInstParent != parentFunc); - } - return false; + auto addrInstParent = getParentFunc(root); + return (addrInstParent != parentFunc); } struct GenericChildrenMigrationContextImpl diff --git a/source/slang/slang.cpp b/source/slang/slang.cpp index a1c8f9b03f..f2598786ff 100644 --- a/source/slang/slang.cpp +++ b/source/slang/slang.cpp @@ -125,6 +125,14 @@ namespace Slang { const char* getBuildTagString() { + if (UnownedStringSlice(SLANG_TAG_VERSION) == "unknown") + { + // If the tag is unknown, then we will try to get the timestamp of the shared library + // and use that as the version string, so that we can at least return something + // that uniquely identifies the build. + static String timeStampString = String(SharedLibraryUtils::getSharedLibraryTimestamp((void*)spCreateSession)); + return timeStampString.getBuffer(); + } return SLANG_TAG_VERSION; }