Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/ltc_examples.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ Refer to the main documentation [here](ltc_backend.md).
```python
import torch
import torch._lazy
import torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND as lazy_backend
import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend

# Register the example LTC backend.
lazy_backend._initialize()
Expand Down
4 changes: 2 additions & 2 deletions examples/ltc_backend_bert.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,7 @@ def main(device='lazy', full_size=False):
losses = train(model, num_epochs, num_training_steps, train_dataloader, device)

# Get debug information from LTC
if 'torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND' in sys.modules:
if 'torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND' in sys.modules:
computation = lazy_backend.get_latest_computation()
if computation:
print(computation.debug_string())
Expand Down Expand Up @@ -148,7 +148,7 @@ def main(device='lazy', full_size=False):
torch._lazy.ts_backend.init()

elif args.device == "MLIR_EXAMPLE":
import torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND as lazy_backend
import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend

lazy_backend._initialize()

Expand Down
4 changes: 2 additions & 2 deletions examples/ltc_backend_mnist.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def forward(self, x):
torch._lazy.mark_step()

# Get debug information from LTC
if 'torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND' in sys.modules:
if 'torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND' in sys.modules:
computation = lazy_backend.get_latest_computation()
if computation:
print(computation.debug_string())
Expand Down Expand Up @@ -93,7 +93,7 @@ def forward(self, x):
torch._lazy.ts_backend.init()

elif args.device == "MLIR_EXAMPLE":
import torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND as lazy_backend
import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend

lazy_backend._initialize()

Expand Down
2 changes: 1 addition & 1 deletion python/test/lazy_backend/device_data_name.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import torch
import torch._lazy

import torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND as lazy_backend
import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend

from run_test import run_test

Expand Down
2 changes: 1 addition & 1 deletion python/torch_mlir/csrc/base_lazy_backend/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ target_link_libraries(torch_mlir_ltc_backend

message(STATUS "TORCH_CXXFLAGS=${TORCH_CXXFLAGS} -Wno-pedantic")
set_target_properties(torch_mlir_ltc_backend PROPERTIES
LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/base_lazy_backend"
LIBRARY_OUTPUT_DIRECTORY "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs"
OUTPUT_NAME lib_torch_mlir_ltc
PREFIX ""
SUFFIX ".so"
Expand Down
6 changes: 5 additions & 1 deletion python/torch_mlir/csrc/reference_lazy_backend/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ mlir_configure_python_dev_packages()
# Library definition
###########################################################################

set(LIBRARY_OUTPUT_PATH "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/reference_lazy_backend")
set(LIBRARY_OUTPUT_PATH "${TORCH_MLIR_PYTHON_PACKAGES_DIR}/torch_mlir/torch_mlir/_mlir_libs")
set(OUTPUT_NAME "_REFERENCE_LAZY_BACKEND")

if(TORCH_MLIR_ENABLE_LTC)
Expand Down Expand Up @@ -64,6 +64,10 @@ if(TORCH_MLIR_ENABLE_LTC)
CXX_VISIBILITY_PRESET "hidden"
COMPILE_FLAGS "${TORCH_CXXFLAGS} -Wno-pedantic"
)
mlir_python_setup_extension_rpath(reference_lazy_backend)

torch_mlir_python_target_compile_options(reference_lazy_backend)
mlir_check_all_link_libraries(reference_lazy_backend)
else()
# To avoid import errors when LTC is disabled (and a bunch of checks
# associated with that), we will generate a dummy placeholder library.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
# Also available under a BSD-style license. See LICENSE.

import torch_mlir.reference_lazy_backend._REFERENCE_LAZY_BACKEND as lazy_backend
import torch_mlir._mlir_libs._REFERENCE_LAZY_BACKEND as lazy_backend
import torch
from torch.utils._pytree import tree_map

Expand Down