[SR-7706] Compiler crashes when sending a scalar value from host to device

[swift-ci]

Previous ID	SR-7706
Radar	None
Original Reporter	hongm (JIRA User)
Type	Bug
Status	Resolved
Resolution	Cannot Reproduce

Additional Detail from JIRA

| | | |------------------|-----------------| |Votes | 0 | |Component/s | Swift for TensorFlow | |Labels | Bug | |Assignee | hongm (JIRA) | |Priority | Medium | md5: e028e4d4361405bfd43bb77aa1056d26

Issue Description:

Example snippet:

var batchSize = Int32(10)
let imageSize = Int32(3)
let images = Tensor<Float>(randomUniform: [batchSize, imageSize])
let batchSize = Tensor<Float>(Float(images.shape[0]))
...

This cases a host->device tensor send on images.shape[0]. In google internal S4TF compiler, the SIL code that produces images.shape[0] looks like:

// function_ref specialized Array.subscript.getter
%267 = function_ref @$SSayxSicigs5Int32V_Tg5 : $@convention(method) (Int, @guaranteed Array\<Int32>) -> Int32 // user: %268
%268 = apply %267(%266, %254) : $@convention(method) (Int, @guaranteed Array\<Int32>) -> Int32 // user: %271
strong_release %264 : $_ContiguousArrayStorageBase // id: %269
strong_release %264 : $_ContiguousArrayStorageBase // id: %270
%271 = struct_extract %268 : $Int32, #Int32._value // user: %274

Where %271 is passed into createHostSend(): https://github.com/apple/swift/blob/2f23c822b5742ea5d1813e7be22b7dba7c966404/lib/SILOptimizer/Mandatory/TFPartition.cpp#L2438

However in open source (OSS) S4TF compiler, the SIL code looks like:

%303 = struct_extract %294 : $Array\<Int32>, #Array._buffer // user: %304
%304 = struct_extract %303 : $_ContiguousArrayBuffer\<Int32>, #_ContiguousArrayBuffer._storage // users: %346, %347, %383, %385, %390, %391, %396, %307, %308, %342, %309, %310, %340, %306, %398, %397, %305
%310 = ref_element_addr %304 : $_ContiguousArrayStorageBase, #_ContiguousArrayStorageBase.countAndCapacity // user: %311
%311 = struct_element_addr %310 : $*_ArrayBody, #_ArrayBody._storage // user: %312
%312 = struct_element_addr %311 : $*_SwiftArrayBodyStorage, #_SwiftArrayBodyStorage.count // user: %313
%313 = struct_element_addr %312 : $*Int, #Int._value // users: %348, %314
%314 = load %313 : $*Builtin.Int64 // users: %339, %325

...

%391 = ref_tail_addr %304 : $_ContiguousArrayStorageBase, $Int32 // user: %392
%392 = struct_element_addr %391 : $*Int32, #Int32._value // user: %393
%393 = load %392 : $*Builtin.Int32 // user: %399

Here %393 represents the int literal value of images.shape[0], and that crashes the code at the referenced point of createHostSend() above.

Is this a case of TFDeabstraction pass not turning load/stores into SSA form? Need more investigation.

[rxwei]

hongm (JIRA User) Are these both Linux builds?

[swift-ci]

Comment by Mingsheng Hong (JIRA)

Yes

[swift-ci]

Comment by Mingsheng Hong (JIRA)

Closing as it's no longer reproducible.