diff --git a/CMakeLists.txt b/CMakeLists.txt index 140a082..0f3a61c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,7 +21,6 @@ option(BUILD_BENCHMARK "Build benchmarks in 'benchmark' directory" ON) option(BUILD_DOCUMENTATION "Build the documentation which is online available" OFF) option(BUILD_EXAMPLE "Build the examples" OFF) -set(IPEX_VERSION "2.1.30" CACHE STRING "The version of Intel Extension for PyTorch") set(TARGET_DEVICE "BMG" CACHE STRING "Valid values {PVC, ACM, BMG}") set(ACM_DEVICES ACM) @@ -32,6 +31,7 @@ string( TOUPPER ${TARGET_DEVICE} TARGET_DEVICE_UPPERCASE) if (TARGET_DEVICE_UPPERCASE IN_LIST XE_DEVICES) message(STATUS "Building for PVC/Xe2") add_definitions(-DTARGET_DEVICE=0) + elseif (TARGET_DEVICE_UPPERCASE IN_LIST ACM_DEVICES) message(STATUS "Building for Alchemist") add_definitions(-DTARGET_DEVICE=1) @@ -39,69 +39,6 @@ else() message(ERROR "Please set a valid device by adding -DTARGET_DEVICE=, where is one of \"PVC\", \"BMG\", or \"ACM\"") endif() -message(WARNING "cmake build system is still WIP.") - -# Check if either BUILD_PYBIND or BUILD_TORCH_TEST is ON -if(BUILD_PYBIND OR BUILD_TORCH_TEST OR BUILD_TORCH_BENCHMARK) - # Check if libtorch directory exists - if(NOT EXISTS "${PROJECT_SOURCE_DIR}/extern/libtorch") - # Notify that we are downloading the necessary libtorch files - message(STATUS "libtorch not found in ${PROJECT_SOURCE_DIR}/extern/. Downloading...") - - # Determine the appropriate filenames based on the IPEX_VERSION - if(IPEX_VERSION VERSION_EQUAL "2.1.30") - set(LIBTORCH_FILENAME "libtorch-cxx11-abi-shared-with-deps-2.1.0.post2.zip") - set(LIBINTEL_EXT_FILENAME "libintel-ext-pt-cxx11-abi-2.1.30.post0.run") - elseif(IPEX_VERSION VERSION_EQUAL "2.1.20") - set(LIBTORCH_FILENAME "libtorch-cxx11-abi-shared-with-deps-2.1.0.post0.zip") - set(LIBINTEL_EXT_FILENAME "libintel-ext-pt-cxx11-abi-2.1.20+xpu.run") - elseif(IPEX_VERSION VERSION_EQUAL "2.1.10") - set(LIBTORCH_FILENAME "libtorch-cxx11-abi-shared-with-deps-2.1.0a0.zip") - set(LIBINTEL_EXT_FILENAME "libintel-ext-pt-cxx11-abi-2.1.10+xpu.run") - else() - message(FATAL_ERROR "Unsupported IPEX_VERSION: ${IPEX_VERSION}") - endif() - - # Set the URLs for the files to be downloaded using the determined filenames - set(LIBTORCH_URL "https://intel-extension-for-pytorch.s3.amazonaws.com/libipex/xpu/${LIBTORCH_FILENAME}") - set(LIBINTEL_EXT_URL "https://intel-extension-for-pytorch.s3.amazonaws.com/libipex/xpu/${LIBINTEL_EXT_FILENAME}") - - # Download libtorch - file(DOWNLOAD "${LIBTORCH_URL}" "${CMAKE_CURRENT_BINARY_DIR}/${LIBTORCH_FILENAME}" - STATUS download_status_libtorch) - list(GET download_status_libtorch 0 download_result_libtorch) - if(NOT download_result_libtorch EQUAL 0) - message(FATAL_ERROR "Failed to download libtorch. Please download it manually with:\nwget ${LIBTORCH_URL}\nand follow install instructions for cppsdk: https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu&version=v2.1.10%2Bxpu") - endif() - - # Download libintel extension - file(DOWNLOAD "${LIBINTEL_EXT_URL}" "${CMAKE_CURRENT_BINARY_DIR}/${LIBINTEL_EXT_FILENAME}" - STATUS download_status_libintel_ext) - list(GET download_status_libintel_ext 0 download_result_libintel_ext) - if(NOT download_result_libintel_ext EQUAL 0) - message(FATAL_ERROR "Failed to download libintel extension. Please download it manually with:\nwget ${LIBINTEL_EXT_URL}\nand follow install instructions for cppsdk: https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu&version=v2.1.10%2Bxpu") - endif() - - # Unzip the downloaded libtorch file - execute_process( - COMMAND ${CMAKE_COMMAND} -E tar xvf "${CMAKE_CURRENT_BINARY_DIR}/${LIBTORCH_FILENAME}" - WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}/extern/" - ) - - # Execute the installer for the libintel extension - execute_process( - COMMAND bash "${CMAKE_CURRENT_BINARY_DIR}/${LIBINTEL_EXT_FILENAME}" install "${PROJECT_SOURCE_DIR}/extern/libtorch/" - WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}/extern/" - ) - - # After unpacking and installation, clean up the downloaded files if needed. - file(REMOVE "${CMAKE_CURRENT_BINARY_DIR}/${LIBTORCH_FILENAME}") - file(REMOVE "${CMAKE_CURRENT_BINARY_DIR}/${LIBINTEL_EXT_FILENAME}") - else() - message(STATUS "libtorch already exists in ${PROJECT_SOURCE_DIR}/extern/.") - endif() - -endif() add_subdirectory(source) diff --git a/README.md b/README.md index 0afdb3f..ab39eee 100644 --- a/README.md +++ b/README.md @@ -108,13 +108,16 @@ conda create -n tiny-dpcpp-nn python=3.10 -y conda activate tiny-dpcpp-nn ``` -Install the latest [ipex](https://intel.github.io/intel-extension-for-pytorch/index.html#installation) via +[Install PyTorch (at the point of writing 2.7)](https://pytorch.org/docs/stable/notes/get_start_xpu.html): -```bash -python -m pip install torch==2.1.0.post2 torchvision==0.16.0.post2 torchaudio==2.1.0.post2 intel-extension-for-pytorch==2.1.30+xpu oneccl_bind_pt==2.1.300+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ ``` +pip install -r requirements.txt +``` + +Verify that drivers and PyTorch are installed correctly: -*Note* please ensure that the IPEX version (2.1.30 in this example) is the same as used in `IPEX_VERSION` in `tiny-dpcpp-nn/CMakeLists.txt` +``` +python -c "import torch; print(torch.xpu.is_available())" Install the module (if no `TARGET_DEVICE` is set, the target_device in setup.py is set to `ARC`. Currently `PVC` and `ARC` is supported): ```bash @@ -122,10 +125,6 @@ cd dpcpp_bindings TARGET_DEVICE=ARC pip install -e . ``` -Finally, to test the sample scripts and tests, install the requirements: -```bash -cd python && pip install -r requirements.txt -``` ### Test the install To test that the installation was successful, you can do the following four tests. @@ -152,7 +151,7 @@ To have all tests, run: cmake -DTARGET_DEVICE="PVC" -DBUILD_TORCH_TEST="ON" .. ``` -After all tests are build into `build/`, you can run `cd build/ && make tests` to verfiy that the setup is correct. Please note that we provide tests for both the core `dpcpp` implementation and the `libtorch` wrapper implementation. +After all tests are build into `build/`, you cvan run `cd build/ && make tests` to verfiy that the setup is correct. Please note that we provide tests for both the core `dpcpp` implementation and the `libtorch` wrapper implementation. To test whether the pytorch bindings were installed correctly, please run diff --git a/dpcpp_bindings/CMakeLists.txt b/dpcpp_bindings/CMakeLists.txt index 7e2324a..65379c4 100644 --- a/dpcpp_bindings/CMakeLists.txt +++ b/dpcpp_bindings/CMakeLists.txt @@ -4,15 +4,11 @@ # Include headers find_package(IntelSYCL REQUIRED) -find_package(Torch REQUIRED PATHS ${TORCH_CMAKE_PREFIX_PATH}) +set(CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX_PATH}" CACHE PATH "Path to Torch") +find_package(Torch REQUIRED) # Add include directories -set(LIBTORCH_DIR "${PROJECT_SOURCE_DIR}/extern/libtorch/share/cmake/") -list(APPEND CMAKE_PREFIX_PATH "${LIBTORCH_DIR}") - -find_package(IPEX REQUIRED) - -include_directories(${TORCH_INCLUDE_DIRS}) +# include_directories(${TORCH_INCLUDE_DIRS}) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/../include) # Set source files for the Python bindings set(PYBIND_SOURCES @@ -24,7 +20,7 @@ pybind11_add_module(tiny_dpcpp_nn_pybind_module ${PYBIND_SOURCES}) find_library(TORCH_PYTHON_LIBRARY torch_python PATHS "${TORCH_CMAKE_PREFIX_PATH}/../../lib" REQUIRED) # Link any required libraries -target_link_libraries(tiny_dpcpp_nn_pybind_module PRIVATE network ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY} ${TORCH_IPEX_LIBRARIES}) +target_link_libraries(tiny_dpcpp_nn_pybind_module PRIVATE network ${TORCH_LIBRARIES} ${TORCH_PYTHON_LIBRARY}) target_include_directories(tiny_dpcpp_nn_pybind_module PRIVATE "${PROJECT_SOURCE_DIR}/extern/json" "${PROJECT_SOURCE_DIR}/extern/pybind11_json") add_sycl_to_target(TARGET tiny_dpcpp_nn_pybind_module SOURCES ${PYBIND_SOURCES}) diff --git a/dpcpp_bindings/pyproject.toml b/dpcpp_bindings/pyproject.toml new file mode 100644 index 0000000..e909578 --- /dev/null +++ b/dpcpp_bindings/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools>=62.6.0", "wheel", "torch", "ninja"] +build-backend = "setuptools.build_meta" diff --git a/dpcpp_bindings/requirements.txt b/dpcpp_bindings/requirements.txt new file mode 100644 index 0000000..61c3fc2 --- /dev/null +++ b/dpcpp_bindings/requirements.txt @@ -0,0 +1,8 @@ +--index-url https://download.pytorch.org/whl/xpu +--extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ +torch==2.7.0 +torchvision==0.22.0 +torchaudio==2.7.0 +ninja==1.11.1.1 +# intel-extension-for-pytorch==2.7.10+xpu +# oneccl_bind_pt==2.7.0+xpu \ No newline at end of file diff --git a/dpcpp_bindings/setup.py b/dpcpp_bindings/setup.py index 1f75811..0b7b444 100644 --- a/dpcpp_bindings/setup.py +++ b/dpcpp_bindings/setup.py @@ -45,7 +45,7 @@ def build_extension(self, ext: CMakeExtension) -> None: # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code # from Python. - target_device = os.environ.get("TARGET_DEVICE", "ARC") + target_device = os.environ.get("TARGET_DEVICE", "PVC") # Validate TARGET_DEVICE if target_device not in {"ARC", "PVC"}: @@ -73,49 +73,15 @@ def build_extension(self, ext: CMakeExtension) -> None: # In this example, we pass in the version to C++. You might not need to. cmake_args += [f"-DEXAMPLE_VERSION_INFO={self.distribution.get_version()}"] - if self.compiler.compiler_type != "msvc": - # Using Ninja-build since it a) is available as a wheel and b) - # multithreads automatically. MSVC would require all variables be - # exported for Ninja to pick it up, which is a little tricky to do. - # Users can override the generator with CMAKE_GENERATOR in CMake - # 3.15+. - if not cmake_generator or cmake_generator == "Ninja": - try: - import ninja - - ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" - cmake_args += [ - "-GNinja", - f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", - ] - except ImportError: - pass - - else: - # Single config generators are handled "normally" - single_config = any(x in cmake_generator for x in {"NMake", "Ninja"}) - - # CMake allows an arch-in-generator style for backward compatibility - contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"}) - - # Specify the arch if using MSVC generator, but only if it doesn't - # contain a backward-compatibility arch spec already in the - # generator name. - if not single_config and not contains_arch: - cmake_args += ["-A", PLAT_TO_CMAKE[self.plat_name]] - - # Multi-config generators have a different way to specify configs - if not single_config: - cmake_args += [ - f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{cfg.upper()}={extdir}" - ] - build_args += ["--config", cfg] - - if sys.platform.startswith("darwin"): - # Cross-compile support for macOS - respect ARCHFLAGS if set - archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", "")) - if archs: - cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))] + # Using ninja + if not cmake_generator or cmake_generator == "Ninja": + import ninja + + ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" + cmake_args += [ + "-GNinja", + f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", + ] # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level # across all generators. diff --git a/dpcpp_bindings/tiny_dpcpp_nn/modules.py b/dpcpp_bindings/tiny_dpcpp_nn/modules.py index 8d2bf69..f3c2ae0 100644 --- a/dpcpp_bindings/tiny_dpcpp_nn/modules.py +++ b/dpcpp_bindings/tiny_dpcpp_nn/modules.py @@ -3,7 +3,7 @@ import numpy as np import time -from tiny_dpcpp_nn_pybind_module import ( +from .tiny_dpcpp_nn_pybind_module import ( Activation, create_network, create_encoding, diff --git a/include/common/tensor_legacy_helper.h b/include/common/tensor_legacy_helper.h new file mode 100644 index 0000000..2abfeb9 --- /dev/null +++ b/include/common/tensor_legacy_helper.h @@ -0,0 +1,18 @@ +#include + +namespace xpu { +namespace dpcpp { + + torch::Tensor fromUSM(void *ptr, torch::Dtype dtype, std::vector sizes, bool clone = true) { + const torch::TensorOptions &options = + torch::TensorOptions().dtype(dtype).device(torch::kXPU); + if (clone) { + return torch::from_blob(ptr, sizes, options).clone(); + } + else { + return torch::from_blob(ptr, sizes, options); + } + } + +} // namespace dpcpp +} // namespace xpu diff --git a/include/encodings/frequency.h b/include/encodings/frequency.h index 81dc9e2..e0e3ccf 100644 --- a/include/encodings/frequency.h +++ b/include/encodings/frequency.h @@ -121,7 +121,7 @@ class FrequencyEncoding : public Encoding { const uint32_t outputs_per_input = n_frequencies * 2; float result = 0.0f; - for (int k = 0; k < outputs_per_input; ++k) { + for (int k = 0; k < int(outputs_per_input); ++k) { result += (float)dL_doutput(i, j * outputs_per_input + k) * dy_dx[i * n_dims_to_encode * outputs_per_input + j * outputs_per_input + k]; } loc_dL_dinput(i, j) = result; diff --git a/include/encodings/grid.h b/include/encodings/grid.h index aa2ba4d..b7e9909 100644 --- a/include/encodings/grid.h +++ b/include/encodings/grid.h @@ -258,7 +258,7 @@ void kernel_grid_backward_input(const size_t num_elements, dL_dx(i, dim) = 0.0f; } - for (int k = 0; k < num_grid_features; ++k) { + for (int k = 0; k < int(num_grid_features); ++k) { const float dL_dy_local = (float)dL_dy_rm(i, k); for (uint32_t dim = 0; dim < N_POS_DIMS; ++dim) { diff --git a/include/network/Network.h b/include/network/Network.h index 45c079c..c58c6df 100644 --- a/include/network/Network.h +++ b/include/network/Network.h @@ -63,15 +63,19 @@ template class Network : public NetworkBase { Network(sycl::queue &q, const int n_hidden_layers, const int input_width, const int network_width, const int output_width, const WeightInitMode mode, const bool use_bias) - : m_q(q), input_width_(PadWidths(input_width, network_width)), - output_width_(PadWidths(output_width, network_width)), original_input_width_(input_width), - original_output_width_(output_width), n_hidden_layers_(NonNegative(n_hidden_layers)), - network_width_(NonNegative(network_width)), - m_weights_matrices(get_n_hidden_layers() + 1, get_input_width(), get_network_width(), get_network_width(), + : m_q(q), + input_width_(PadWidths(input_width, network_width)), + output_width_(PadWidths(output_width, network_width)), + original_output_width_(output_width), + original_input_width_(input_width), + n_hidden_layers_(NonNegative(n_hidden_layers)), + network_width_(NonNegative(network_width)), + use_bias_(use_bias), + m_weights_matrices(get_n_hidden_layers() + 1, get_input_width(), get_network_width(), get_network_width(), get_network_width(), get_network_width(), get_output_width(), m_q), - m_weightsT_matrices(get_n_hidden_layers() + 1, get_network_width(), get_input_width(), get_network_width(), - get_network_width(), get_output_width(), get_network_width(), m_q), - use_bias_(use_bias) { + m_weightsT_matrices(get_n_hidden_layers() + 1, get_network_width(), get_input_width(), get_network_width(), + get_network_width(), get_output_width(), get_network_width(), m_q) + { SanityCheck(); initialize_weights_matrices(input_width, output_width, mode); @@ -135,14 +139,14 @@ template class Network : public NetworkBase { throw std::runtime_error(errorMessage); } - if (get_input_width() > network_width_) { + if (int(get_input_width()) > network_width_) { std::string errorMessage = "Input width of " + std::to_string(get_input_width()) + " is not supported. Value must be <= network width (" + std::to_string(network_width_) + ")."; throw std::runtime_error(errorMessage); } - if (get_output_width() > network_width_) { + if (int(get_output_width()) > network_width_) { std::string errorMessage = "Input width of " + std::to_string(get_output_width()) + " is not supported. Value must be <= network width (" + std::to_string(network_width_) + ")."; @@ -158,14 +162,14 @@ template class Network : public NetworkBase { if (network_width_ != 16 && network_width_ != 32 && network_width_ != 64 && network_width_ != 128) throw std::invalid_argument("Network width has to be a power of 2 between 16 and 128."); - if (network_width_ != get_input_width() || network_width_ != get_output_width()) + if (network_width_ != int(get_input_width()) || network_width_ != int(get_output_width())) throw std::invalid_argument("Only networks with same input, layer and output width are allowed."); } ///@brief Helper function which sets values of the weights matrices to 0 if /// the actual input/output width was padded to the network-allowed input/output width. void ZeroWeightsPadding(const int unpadded_input_width, const int unpadded_output_width) { - if (unpadded_input_width > get_input_width() || unpadded_output_width > get_output_width()) + if (unpadded_input_width > int(get_input_width()) || unpadded_output_width > int(get_output_width())) throw std::invalid_argument("Padded weights width cannot be less than the unpadded."); /// we need to copy everything here since we do not want to have an implicit copy of 'this' @@ -213,7 +217,7 @@ template class Network : public NetworkBase { } // output matrix set columns to 0 - const int output_matrix_pos = m_weights_matrices.GetNumberOfMatrices() - 1; + // const int output_matrix_pos = m_weights_matrices.GetNumberOfMatrices() - 1; if (unpadded_output_width != padded_output_width) { DeviceMatrixView weights = m_weights_matrices.Back(); m_q.parallel_for( diff --git a/include/network/SwiftNetMLP.h b/include/network/SwiftNetMLP.h index 4c5fb49..28ef154 100644 --- a/include/network/SwiftNetMLP.h +++ b/include/network/SwiftNetMLP.h @@ -12,7 +12,7 @@ #pragma once -#include +#include #include #include #include @@ -250,6 +250,7 @@ template class SwiftNetMLP : public Network { private: /// Generate the relevant kernel class. Has to be called in constructor void checkKernels() const { + const auto arch = Network::get_queue().get_device().template get_info(); switch (arch) { case syclex::architecture::intel_gpu_pvc: @@ -332,7 +333,7 @@ template class SwiftNetMLP : public Network { std::to_string(intermediate_output_forward.output_n()) + ")"); } - if (intermediate_output_forward.GetNumberOfMatrices() != Network::get_n_hidden_layers() + 2) { + if (int(intermediate_output_forward.GetNumberOfMatrices()) != Network::get_n_hidden_layers() + 2) { throw std::invalid_argument("Not enough matrices in intermediate_output_forward array, expected: " + std::to_string(Network::get_n_hidden_layers() + 2) + " but got: " + std::to_string(intermediate_output_forward.GetNumberOfMatrices())); @@ -369,7 +370,7 @@ template class SwiftNetMLP : public Network { std::to_string(intermediate_output_forward.middle_n()) + ", " + std::to_string(intermediate_output_forward.output_n()) + ")"); - if (intermediate_output_forward.GetNumberOfMatrices() != Network::get_n_hidden_layers() + 2) + if (int(intermediate_output_forward.GetNumberOfMatrices()) != Network::get_n_hidden_layers() + 2) throw std::invalid_argument("Not enough matrices in intermediate_output_forward array. Required: " + std::to_string(Network::get_n_hidden_layers() + 2) + ", available: " + std::to_string(intermediate_output_forward.GetNumberOfMatrices())); @@ -393,7 +394,7 @@ template class SwiftNetMLP : public Network { std::to_string(intermediate_output_backward.middle_n()) + ", " + std::to_string(intermediate_output_backward.output_n()) + ")"); - if (intermediate_output_backward.GetNumberOfMatrices() != Network::get_n_hidden_layers() + 1) + if (int(intermediate_output_backward.GetNumberOfMatrices()) != Network::get_n_hidden_layers() + 1) throw std::invalid_argument("Not enough matrices in intermediate_output_backward array. Required: " + std::to_string(Network::get_n_hidden_layers() + 1) + ", available: " + std::to_string(intermediate_output_backward.GetNumberOfMatrices())); diff --git a/include/optimizers/adam.h b/include/optimizers/adam.h index 2e1e04f..13399c9 100644 --- a/include/optimizers/adam.h +++ b/include/optimizers/adam.h @@ -29,8 +29,8 @@ class AdamOptimizer : public Optimizer { DeviceMem m_first_moments; DeviceMem m_second_moments; - int m_output_rows; - int m_n_hidden_layers; + // int m_output_rows; + // int m_n_hidden_layers; float m_learning_rate = 1e-3f; float m_l2_reg = 1e-8f; }; diff --git a/include/tnn_api.h b/include/tnn_api.h index 00faecd..3af1ded 100644 --- a/include/tnn_api.h +++ b/include/tnn_api.h @@ -12,6 +12,7 @@ #pragma once +#include #include #include @@ -32,6 +33,7 @@ #include "json.hpp" #include "oneapi/mkl.hpp" #include "result_check.h" +#include "tensor_legacy_helper.h" extern template class SwiftNetMLP; extern template class SwiftNetMLP; @@ -292,11 +294,11 @@ class Module { private: sycl::queue &sycl_queue_; static sycl::queue &InitQueue() { - // return c10::xpu::getCurrentXPUStream().queue(); - auto device_type = c10::DeviceType::XPU; - c10::impl::VirtualGuardImpl impl(device_type); - c10::Stream xpu_stream = impl.getStream(c10::Device(device_type)); - return xpu::get_queue_from_stream(xpu_stream); + return c10::xpu::getCurrentXPUStream().queue(); + // auto device_type = c10::DeviceType::XPU; + // c10::impl::VirtualGuardImpl impl(device_type); + // c10::Stream xpu_stream = impl.getStream(c10::Device(device_type)); + // return xpu::get_queue_from_stream(xpu_stream); } }; @@ -381,7 +383,7 @@ class EncodingModule : public Module { } std::unique_ptr model_ctx = nullptr; - const size_t batch_size = grad_output.sizes()[0]; + const auto batch_size = grad_output.sizes()[0]; const int64_t input_width = encoding_->get_input_width(); // Assuming input_from_fwd should have dimensions [batch_size, diff --git a/source/optimizers/adam.cpp b/source/optimizers/adam.cpp index 9304a8e..984f458 100644 --- a/source/optimizers/adam.cpp +++ b/source/optimizers/adam.cpp @@ -148,16 +148,20 @@ void AdamOptimizer::step(queue q, float loss_scale, DeviceMem &weights, De auto first_moment = m_first_moments.data(); auto second_moment = m_second_moments.data(); + bf16* const weights_data = weights.data(); + bf16* const weightsT_data = weightsT.data(); + bf16* const gradients_data = gradients.data(); + q.parallel_for<>(range<1>(n_elements), [=](id<1> idx) { adam_step(idx, n_elements, relative_weight_decay, absolute_weight_decay, weight_clipping_magnitude, loss_scale, learning_rate, non_matrix_learning_rate_factor, beta1, beta2, epsilon, lower_lr_bound, - upper_lr_bound, l2_reg, weights.data(), gradients.data(), first_moment, second_moment, WIDTH); + upper_lr_bound, l2_reg, weights_data, gradients_data, first_moment, second_moment, WIDTH); }).wait(); q.parallel_for<>(range<1>(n_elements), [=](id<1> idx) { adam_stepT(idx, n_elements, relative_weight_decay, absolute_weight_decay, weight_clipping_magnitude, loss_scale, learning_rate, non_matrix_learning_rate_factor, beta1, beta2, epsilon, lower_lr_bound, - upper_lr_bound, l2_reg, weightsT.data(), gradients.data(), first_moment, second_moment, WIDTH); + upper_lr_bound, l2_reg, weightsT_data, gradients_data, first_moment, second_moment, WIDTH); }).wait(); } diff --git a/source/optimizers/sgd.cpp b/source/optimizers/sgd.cpp index 1a8dbfe..73a8e52 100644 --- a/source/optimizers/sgd.cpp +++ b/source/optimizers/sgd.cpp @@ -137,16 +137,20 @@ void SGDOptimizer::step(queue q, float loss_scale, DeviceMem &weights, Dev const int output_rows = m_output_rows; const int n_hidden_layers = m_n_hidden_layers; + bf16* const weights_data = weights.data(); + bf16* const weightsT_data = weightsT.data(); + bf16* const gradients_data = gradients.data(); + // Perform the SGD update for weight matrices q.parallel_for<>(range<1>(n_elements), [=](id<1> idx) { - sgd_step(idx, n_elements, output_rows, n_hidden_layers, loss_scale, learning_rate, l2_reg, weights.data(), - gradients.data(), WIDTH); + sgd_step(idx, n_elements, output_rows, n_hidden_layers, loss_scale, learning_rate, l2_reg, weights_data, + gradients_data, WIDTH); }).wait(); // Perform the SGD update for transposed weight matrices q.parallel_for<>(range<1>(n_elements), [=](id<1> idx) { - sgd_stepT(idx, n_elements, output_rows, n_hidden_layers, loss_scale, learning_rate, l2_reg, weightsT.data(), - gradients.data(), WIDTH); + sgd_stepT(idx, n_elements, output_rows, n_hidden_layers, loss_scale, learning_rate, l2_reg, weightsT_data, + gradients_data, WIDTH); }).wait(); }