01_texture_inference.comp

/*!
  \file 01_texture_inference.comp
  \author Sho Ikeda
  \brief HLSL compute kernel for texture generation using MLP
  \copyright Copyright (c) 2026 Advanced Micro Devices, Inc. All Rights Reserved.

  SPDX-License-Identifier: MIT
*/

#include "mlp.hlsl"
#include "texture_inference_common.hlsl"

ByteAddressBuffer UvBuffer;
RWByteAddressBuffer OutputBuffer;
ByteAddressBuffer WeightBuffer;
ByteAddressBuffer BiasBuffer;
int TEST_WEIGHT_MATRIX_SIZE_FIRST;
int TEST_WEIGHT_MATRIX_SIZE_HIDDEN;

template <typename Type>
void inference(const uint threadId)
{
  const dx::linalg::ComponentEnum elemType = mininn::impl::TypeTraits<Type>::COMPONENT_TYPE;
  const dx::linalg::MatrixLayoutEnum layout = (dx::linalg::MatrixLayoutEnum)MINIDXNN_WEIGHT_MATRIX_LAYOUT;
  using ActivationHiddenT = mininn:: MINIDXNN_ACTIVATION_HIDDEN_TYPE;
  using ActivationLastT = mininn:: MINIDXNN_ACTIVATION_LAST_TYPE;

  texkernel::inferenceStep<Type, MINIDXNN_NUM_LAYERS, MINIDXNN_HIDDEN_LAYER_DIMENSIONS,
      elemType, layout, ActivationHiddenT, ActivationLastT,
      MINIDXNN_WEIGHT_MATRIX_ALIGNMENT, MINIDXNN_WEIGHT_MATRIX_VECTOR_STRIDE_ALIGNMENT, MINIDXNN_BIAS_VECTOR_ALIGNMENT,
      (MINIDXNN_HAS_BIAS != 0)>(
      threadId, UvBuffer, OutputBuffer, WeightBuffer, BiasBuffer,
      uint2((uint)TEST_WEIGHT_MATRIX_SIZE_FIRST, (uint)TEST_WEIGHT_MATRIX_SIZE_HIDDEN),
      MINIDXNN_NUM_TASKS);
}

[numthreads(MINIDXNN_NUM_THREADS_X, 1, 1)]
void inferenceF32Kernel(const uint3 groupThreadId : SV_GroupThreadID, const uint3 groupId : SV_GroupID)
{
  const uint threadId = groupId.x * MINIDXNN_NUM_THREADS_X + groupThreadId.x;
  inference<float>(threadId);
}

[numthreads(MINIDXNN_NUM_THREADS_X, 1, 1)]
void inferenceF16Kernel(const uint3 groupThreadId : SV_GroupThreadID, const uint3 groupId : SV_GroupID)
{
  const uint threadId = groupId.x * MINIDXNN_NUM_THREADS_X + groupThreadId.x;
  inference<half>(threadId);
}