Skip to content

Instantly share code, notes, and snippets.

@ManfeiBai
Created July 12, 2024 21:16
Show Gist options
  • Save ManfeiBai/ea9cf0f0db409f4bcdd45f6af65c638d to your computer and use it in GitHub Desktop.
Save ManfeiBai/ea9cf0f0db409f4bcdd45f6af65c638d to your computer and use it in GitHub Desktop.
0
(torch310) root@b7b12c30e894:/pytorch/xla# cat test_7665.py
import torch
import torch_xla
import torch_xla.experimental.fori_loop
from torch._higher_order_ops.while_loop import while_loop
import torch_xla.core.xla_model as xm
class TPUComputation:
def __init__(self):
self.device = xm.xla_device()
self.init_x = torch.tensor([1], device=self.device)
self.init_y = torch.tensor([1], device=self.device)
self.init_z = torch.tensor([1], device=self.device)
self.iteri = torch.tensor(10, device=self.device)
self.quantity = torch.tensor(0, device=self.device)
def cond_fn(self, iteri, x, y, z, q=None):
return iteri > 0
def body_fn(self, iteri, x, y, z, q=None):
return iteri - 1, x.clone(), y.add(1), z + self.quantity # Problemmatic Line
def compute(self):
result = while_loop(
self.cond_fn,
self.body_fn,
(self.iteri, self.init_x, self.init_y, self.init_z),
)
return result
if __name__ == "__main__":
computation = TPUComputation()
result = computation.compute()
print(result)
(torch310) root@b7b12c30e894:/pytorch/xla# PJRT_DEVICE=TPU python test_7665.py
body computation: !!!!!!!!!
HloModule PyLoweringContext.21.29, entry_computation_layout={((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}))->(s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0})}
%PyLoweringContext.7 (p0.10: s64[], p1.14: s64[1], p2.17: s64[1], p3.22: s64[1], UnusedArgumentsPlaceholder.26: s64[1]) -> (s64[], s64[], s64[1], s64[1], s64[1]) {
%UnusedArgumentsPlaceholder.26 = s64[1]{0} parameter(4)
%p0.10 = s64[] parameter(0)
%constant.9 = s64[] constant(1)
%constant.8 = s64[] constant(1)
%multiply.11 = s64[] multiply(s64[] %constant.9, s64[] %constant.8)
%subtract.12 = s64[] subtract(s64[] %p0.10, s64[] %multiply.11)
%constant.13 = s64[] constant(0)
%p1.14 = s64[1]{0} parameter(1)
%p2.17 = s64[1]{0} parameter(2)
%constant.16 = s64[] constant(1)
%constant.15 = s64[] constant(1)
%multiply.18 = s64[] multiply(s64[] %constant.16, s64[] %constant.15)
%broadcast.19 = s64[1]{0} broadcast(s64[] %multiply.18), dimensions={}
%add.20 = s64[1]{0} add(s64[1]{0} %p2.17, s64[1]{0} %broadcast.19)
%p3.22 = s64[1]{0} parameter(3)
%constant.21 = s64[] constant(1)
%multiply.23 = s64[] multiply(s64[] %constant.13, s64[] %constant.21)
%broadcast.24 = s64[1]{0} broadcast(s64[] %multiply.23), dimensions={}
%add.25 = s64[1]{0} add(s64[1]{0} %p3.22, s64[1]{0} %broadcast.24)
ROOT %tuple.27 = (s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) tuple(s64[] %subtract.12, s64[] %constant.13, s64[1]{0} %p1.14, s64[1]{0} %add.20, s64[1]{0} %add.25)
}
ENTRY %PyLoweringContext.21.29 (in.1: (s64[], s64[1], s64[1], s64[1], s64[1])) -> (s64[], s64[], s64[1], s64[1], s64[1]) {
%in.1 = (s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) parameter(0)
%get-tuple-element.2 = s64[] get-tuple-element((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=0
%get-tuple-element.3 = s64[1]{0} get-tuple-element((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=1
%get-tuple-element.4 = s64[1]{0} get-tuple-element((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=2
%get-tuple-element.5 = s64[1]{0} get-tuple-element((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=3
%get-tuple-element.6 = s64[1]{0} get-tuple-element((s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=4
ROOT %call.28 = (s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) call(s64[] %get-tuple-element.2, s64[1]{0} %get-tuple-element.3, s64[1]{0} %get-tuple-element.4, s64[1]{0} %get-tuple-element.5, s64[1]{0} %get-tuple-element.6), to_apply=%PyLoweringContext.7
}
cond computation: !!!!!!!!!
HloModule PyLoweringContext.8.16, entry_computation_layout={((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}))->pred[]}
%PyLoweringContext.7 (p0.9: s64[], UnusedArgumentsPlaceholder.11: s64[], UnusedArgumentsPlaceholder.12: s64[1], UnusedArgumentsPlaceholder.13: s64[1], UnusedArgumentsPlaceholder.14: s64[1]) -> pred[] {
%p0.9 = s64[] parameter(0)
%constant.8 = s64[] constant(0)
ROOT %compare.10 = pred[] compare(s64[] %p0.9, s64[] %constant.8), direction=GT
%UnusedArgumentsPlaceholder.11 = s64[] parameter(1)
%UnusedArgumentsPlaceholder.12 = s64[1]{0} parameter(2)
%UnusedArgumentsPlaceholder.13 = s64[1]{0} parameter(3)
%UnusedArgumentsPlaceholder.14 = s64[1]{0} parameter(4)
}
ENTRY %PyLoweringContext.8.16 (in.1: (s64[], s64[], s64[1], s64[1], s64[1])) -> pred[] {
%in.1 = (s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) parameter(0)
%get-tuple-element.2 = s64[] get-tuple-element((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=0
%get-tuple-element.3 = s64[] get-tuple-element((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=1
%get-tuple-element.4 = s64[1]{0} get-tuple-element((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=2
%get-tuple-element.5 = s64[1]{0} get-tuple-element((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=3
%get-tuple-element.6 = s64[1]{0} get-tuple-element((s64[], s64[], s64[1]{0}, s64[1]{0}, s64[1]{0}) %in.1), index=4
ROOT %call.15 = pred[] call(s64[] %get-tuple-element.2, s64[] %get-tuple-element.3, s64[1]{0} %get-tuple-element.4, s64[1]{0} %get-tuple-element.5, s64[1]{0} %get-tuple-element.6), to_apply=%PyLoweringContext.7
}
F0712 21:12:19.562320 1782465 debug_macros.h:21] Non-OK-status: status.status()
Status: INVALID_ARGUMENT: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (in: (s64[], s64[], s64[1], s64[1], s64[1])) -> pred[]; body: (in: (s64[], s64[1], s64[1], s64[1], s64[1])) -> (s64[], s64[], s64[1], s64[1], s64[1]); init: (s64[], s64[], s64[1], s64[1], s64[1])..:
*** Begin stack trace ***
tsl::CurrentStackTrace[abi:cxx11]()
xla::XlaComputation ConsumeValue<xla::XlaComputation>(absl::lts_20230802::StatusOr<xla::XlaComputation>&&)
_PyObject_MakeTpCall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyObject_FastCallDictTstate
_PyObject_Call_Prepend
_PyObject_MakeTpCall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyObject_FastCallDictTstate
_PyObject_Call_Prepend
PyObject_Call
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyObject_FastCallDictTstate
_PyObject_Call_Prepend
_PyObject_MakeTpCall
_PyEval_EvalFrameDefault
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
_PyFunction_Vectorcall
_PyEval_EvalFrameDefault
PyEval_EvalCode
_PyRun_SimpleFileObject
_PyRun_AnyFileObject
Py_RunMain
Py_BytesMain
__libc_start_main
*** End stack trace ***
*** Check failure stack trace: ***
@ 0x7f90c3f910d9 absl::lts_20230802::log_internal::LogMessageFatal::~LogMessageFatal()
@ 0x7f90bbab75e1 ConsumeValue<>()
@ 0x7f90bbb1bf86 torch_xla::(anonymous namespace)::CreateComputation()
@ 0x7f90bbb1c3c4 pybind11::cpp_function::initialize<>()::{lambda()#3}::_FUN()
@ 0x7f90bbae9a76 pybind11::cpp_function::dispatcher()
@ 0x4fdc87 cfunction_call
https://symbolize.stripped_domain/r/?trace=7f91780d4ce1,7f91783d213f,7f90c3f909bc,7f90c3f910d8,7f90bbab75e0,7f90bbb1bf85,7f90bbb1c3c3,7f90bbae9a75,4fdc86&map=
*** SIGABRT received by PID 1782465 (TID 1782465) on cpu 0 from PID 1782465; stack trace: ***
PC: @ 0x7f91780d4ce1 (unknown) raise
@ 0x7f8c83003681 944 (unknown)
@ 0x7f91783d2140 (unknown) (unknown)
@ 0x7f90c3f909bd 32 absl::lts_20230802::log_internal::LogMessage::Die()
@ 0x7f90c3f910d9 16 absl::lts_20230802::log_internal::LogMessageFatal::~LogMessageFatal()
@ 0x7f90bbab75e1 112 ConsumeValue<>()
@ 0x7f90bbb1bf86 672 torch_xla::(anonymous namespace)::CreateComputation()
@ 0x7f90bbb1c3c4 192 pybind11::cpp_function::initialize<>()::{lambda()#3}::_FUN()
@ 0x7f90bbae9a76 704 pybind11::cpp_function::dispatcher()
@ 0x4fdc87 (unknown) cfunction_call
@ ... and at least 2 more frames
https://symbolize.stripped_domain/r/?trace=7f91780d4ce1,7f8c83003680,7f91783d213f,7f90c3f909bc,7f90c3f910d8,7f90bbab75e0,7f90bbb1bf85,7f90bbb1c3c3,7f90bbae9a75,4fdc86&map=
E0712 21:12:19.689653 1782465 coredump_hook.cc:316] RAW: Remote crash data gathering hook invoked.
E0712 21:12:19.689670 1782465 client.cc:269] RAW: Coroner client retries enabled, will retry for up to 30 sec.
E0712 21:12:19.689680 1782465 coredump_hook.cc:411] RAW: Sending fingerprint to remote end.
E0712 21:12:19.689700 1782465 coredump_hook.cc:420] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] stat failed on crash reporting socket /var/google/services/logmanagerd/remote_coredump.socket (Is the listener running?): No such file or directory
E0712 21:12:19.689706 1782465 coredump_hook.cc:472] RAW: Dumping core locally.
E0712 21:13:07.653508 1782465 process_state.cc:805] RAW: Raising signal 6 with default behavior
Aborted (core dumped)
(torch310) root@b7b12c30e894:/pytorch/xla#
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment