iree-org
Implement initial heterogeneous support MVP for CPU (+ maybe something else).
TLDR: programs will be able to selectively allocate buffers that are compatible with all devices they are used on, submit command buffers using resources from each other, and synchronize with semaphores to order operations. Performance is not a goal of this effort - only straight line execution and the addition of passes/APIs to enable the support. An initial MVP will be CPU+CPU only but using nothing any other HAL implementation could do; HIP support is a stretch if anyone is interested in helping out.
Heterogeneous device support builds on the current homogeneous multi-device support by extending what was done there with a few new features (see the sub-issues below).
Once implemented we'll have compiled programs that can run and request allocations on devices selected from a subset of those devices a particular buffer is used on. A policy defined by the application (possibly derived from available devices, but TBD) will be used for selection. For example, a buffer used on CPU and NPU could always be allocated on the NPU if it requires special allocations (vs normal host malloc). Transfers through staging buffers in host memory can be used when two devices cannot import/export though such compiler support may come later.
After improving the compiler the remaining work is entirely within the runtime HAL for better seamless buffer and semaphore interop. For an MVP of CPU + CPU not much is required. For when the stretch goal of CPU + [some accelerator] is used buffers are mostly taken care of as the local executors always map PERSISTENT for use in dispatches via the buffer vtable and as long as the non-CPU device implements mapping properly host pointers can be resolved. Semaphores are the trickier item but the recently added iree_hal_semaphore_import_timepoint/iree_hal_semaphore_export_timepoint can help with that in the mean-time (even if slow to start). Any non-CPU HAL will need to support IREE_HAL_EXTERNAL_TIMEPOINT_TYPE_WAIT_PRIMITIVE to interop with the CPU side.
Improvements to compiler support for eliding more transfers, automatically detecting when transfers are still beneficial (such as in NUMA systems), and declarative topologies for indicating links will be left for future work. Other HAL targets besides CPU and the selected experimental target will also be left for future work. There are several known places with poor performance characteristics around import/export of semaphores and potentially required scoped buffer mapping that will be left for future work. At a baseline any HAL implementation wanting to interop will need to support WAIT_PRIMITIVE and map_range and all other faster paths are considered optimizations.
The largest open design question on the compiler-side is around the required cost modeling we cannot (currently/ever) do. Transfers are inserted today with the assumption that all devices are NUMA and the cost of transferring amortizes the traffic of fetching in consumers across a potentially slow bus without DMA support. To start a coarse non-default flag will be added to elide everything possible. I'm not sure there's any good solution that does not require algorithm-aware analysis that we don't have today (measuring reads/writes and reuse). An explicit attribute on transfer ops inserted in frontends may be required to indicate whether a transfer is load-bearing for performance and should not be elided to allow tuning.
Thanks a lot for taking the time to outline this. We would be happy to help with any issues you created here.
We’ve been doing some prototyping ourselves using the metal/local-task drivers, where we took a shortcut for the “programs will be able to selectively allocate buffers that are compatible with all devices” by simply allocating everything as shared using the metal allocator and manually working around the semaphore issues by hand-modifying the HAL IR.
I have some questions about how to handle buffer import/export—or whether it's even necessary at all for the MVP. My observation so far is that drivers typically require buffers to be allocated using their own allocator, since many of the iree_hal_*_command_buffer_* functions expect a native buffer handle. However, this doesn't seem to be the case for the local-task device, as I couldn’t find any references indicating that it requires its own native handle (which I assume is iree_hal_heap_buffer_t).
Given this, is my understanding correct that for the MVP, buffer import/export might not be required at all? As you mentioned in (#20857/#20855), shared allocations would be performed by the device that supports them (metal, in this example), so the HAL functions on the GPU device would still have access to their native handles, and the local devices wouldn’t care?
If the import/export is still needed would it be something that is handled solely in the runtime or would the transfer ops have to be translated into some form of import/export ops instead of eliding them completely?
If its of any interest for someone, this is the IR snipped i got running for local-task/metal when replacing the local-task allocator with the metal allocator:
metal/local-task example
originated from the tools/test/iree-run-module-multi.mlir testcase, there are not a lot of modifications, mostly to resolve semaphore errors like joining on multiple fences that are not from the same device.
compile command i used:
./build/tools/iree-compile --compile-from=vm --iree-execution-model=async-external --iree-hal-target-device=device_a=local[0] --iree-hal-target-device=device_b=metal --iree-hal-local-target-device-backends=vmvx working_multi_device.mlir -o multi.vmfb
run command:
./build/tools/iree-run-module --module=multi.vmfb --device=local-task --device=metal --input=4xf32=10,11,12,13 --function=multi_device_mul
module attributes {vm.toplevel} {
vm.module public @module {
vm.global.ref private mutable @device_a : !vm.ref<!hal.device>
vm.rodata private @_utf8_hal_device_id_369FCE3B885986F6 {alignment = 1 : i64} "hal.device.id"
vm.rodata private @_utf8_local_2C7344D4E05782F4 {alignment = 1 : i64} "local*"
vm.rodata private @_utf8_hal_executable_format_C7128E2AE1BE720D {alignment = 1 : i64} "hal.executable.format"
vm.rodata private @_utf8_vmvx_bytecode_fb_90AB3AF67FE09641 {alignment = 1 : i64} "vmvx-bytecode-fb"
vm.rodata private @multi_device_mul_dispatch_0_vmvx_bytecode_fb {alignment = 16 : i64, mime_type = "application/x-flatbuffers"} dense<"0x504B03042D000000000000002100BBE603A1FFFFFFFFFFFFFFFF090019006D6F64756C652E6662CAFE010000010010007C040000000000007C040000000000007C04000030000000495245450000000000000000000000000000000000000000000000000000000000000000000000000000000012FCFFFF28000000010000008C010000A0010000A4010000A8010000F402000000000F000403000010000000060000006D6F64756C65000054FCFFFF0C00000010000000200000000400000000020201040000000000000004010000F0000000E0000000010000000400000096FCFFFF08000000380000002D0000006D756C74695F6465766963655F6D756C5F64697370617463685F305F656C656D656E74776973655F345F66333200000012000000010000000000000008000000000000000B000000000000000E00000000000000190000000000000024000000000000002B00000001000000360000000300000041000000030000005100000003000000580000000300000069000000030000007100000003000000790000000300000081000000030000008900000003000000900000000300000099000000030000006AFDFFFF020000000100000068FDFFFF1C000000210000001200000078FDFFFF0C000000180000000100000025000000746F6F6C732F746573742F697265652D72756E2D6D6F64756C652D6D756C74692E6D6C697200000005000000880000007400000060000000440000001C000000010000008000000001000000B400000001000000FC000000EAFDFFFF040000001400000021766D2E6C6973743C21766D2E6275666665723E000000000EFEFFFF040000000A00000021766D2E627566666572000026FEFFFF04000000030000006936340036FEFFFF04000000030000006933320046FEFFFF04000000030000006633320056FEFFFF040000002D0000006D756C74695F6465766963655F6D756C5F64697370617463685F305F656C656D656E74776973655F345F6633320000009AFEFFFF08000000180000000F000000307272726969696969696969695F76000100000004000000C2FEFFFF08000000180000000D00000062696E64696E675F636F756E740000000100000032000000F4FEFFFF0400000010000000000000000000803F00000040000040400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001000000000000009C000000010000000000040009000300A0000000790D0100000000000C01000E02000F040000000000000004000F01000000000000000600110000000000801A028001000300000001801A02C0000003000000028056500000000200020000000300010079500000040002005702006800000000000098000000000079E033018000000200E033008000000300E00C020003000200E0340280000002002A0000060000005650000000000000795A0000000000000A00080000000000040008000C00040008000600080004000A001000040008000C0024002C0004000000080000000C0000000000100014001800000000001C0020002400280000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000504B01021E032D000000000000002100BBE603A1FFFFFFFFFFFFFFFF09001C0000000000000000000000FFFFFFFF6D6F64756C652E6662010018007C040000000000007C040000000000000000000000000000504B06062C000000000000002D002D000000000000000000010000000000000001000000000000005300000000000000C004000000000000504B060700000000130500000000000001000000504B0506000000000100010053000000FFFFFFFF0000"> : vector<1397xi8>
vm.rodata private @_utf8_metal_FB846B13C56E663B {alignment = 1 : i64} "metal"
vm.rodata private @_utf8_metal_msl_fb_BD2413475D7AB9A0 {alignment = 1 : i64} "metal-msl-fb"
vm.rodata private @multi_device_mul_dispatch_1_metal_msl_fb {alignment = 16 : i64, mime_type = "application/x-flatbuffers"} dense<"0x0C0000004D544C31000000006CE7FFFF08000000E8000000010000000400000078E7FFFF2C000000200000000100000001000000080000004C00000002000000010000000000000000000000000000002D0000006D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F663332000000E4E7FFFF4000000004000000F0E7FFFF080000001800000025000000746F6F6C732F746573742F697265652D72756E2D6D6F64756C652D6D756C74692E6D6C69720000002D0000006D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F663332000000010000000400000064E8FFFFE010000004000000D31000004D544C4201800200070000810F000000D3100000000000005800000000000000A300000000000000030100000000000008000000000000000B0100000000000008000000000000001301000000000000C00F00000000000001000000A30000004E414D452E006D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F6633320054595045010002484153482000F7309086295136DEE9D43B0907FF276E6B4B22979D1AA2C636E7F4D08281DC014D44535A0800C00F0000000000004F46465418000000000000000000000000000000000000000000000000005645525308000200070003000200454E4454454E445404000000454E445404000000454E4454DEC0170B0000000014000000A80F0000FFFFFFFF4243C0DE3514000003000000620C3024801005C814000000210C0000A00300000B02210002000000160000000781239141C80449061032399201840C250508191E048B628014450242920B42A41032143808184B0A3252884870C421234412878C1041920264C808B1142043468820C901325284182A282A90317CB05C9120C5C800000089200000240000003222480920624600212B249814212524981419270C85A4906052645C20246582C09A015002809534459430F9D253B98D04148416450096823088200823001A003091A6881226FFD7535006118C003D698A2861F2A5A7EA08498890677222969836AA2888884EC23983088830470006C3088362069216B1810015507304A03047104C010C221CC22002100C220002000051180000640000001BEC23F8FFFFFFFF013005C00F003800FE00908870800778908777C0873630877A708771680373808736688770A0077400E8411EEAA11C00C21CE4211CDAA11CDA001EDE211DDC811ECA411E807060077600887648077768037628877308077668037B288771A08777908736B8877420077A4007200FEC0018DCE11DDA401CEA211DD8811ED2C11DE60120DCE11DDA201DDCC11CE6A10DCC011EDAA01DC2811ED001A00779A88772000877788736708770708779680373808736688770A0077400E8411EEAA11C00C21DDEA10DE6211DCEC11DCA811CDA401FCA411EDE611EDAC01CE0A10DDA211CE8011D007A90877A28078070877768037A908770800778480777388736688770A0077400E8411EEAA11C00621EE8211CC6611DDA001EE4E11DE8A11CC6811EDE411EDA401CEAC11CCCA11CE4A10DE6211DF4A11C003C00887A70877908077328873630077868837608077A4007801EE4A11ECA01D8300642002C002900551020411A404338A4833CB48138D48339984339C8431BB8C33BB44138B0433A84C33C001B8841004801004918000002000000138640180000000013B070480779B0033A688370800778608772688376088771788779C08738A8033780033780830DB7510E6D000F7A600774A0077640077A600774D006E910077A80077A80076D900E78A00778A00778D006E9100776A0077160077A100776D006E9300772A0077320077A300772D006E9600774A0077640077A600774D006E6300772A0077320077A300772D006E6600774A0077640077A600774D006F6100776A0077160077A100776D006F6200774A0077320077A300772D006F6300772A0077320077A300772D006F6400778A0077640077A600774D006F6600774A0077640077A600774D006F6900776A00771200778A00771200778D006F610077280077A10077280077A10077280076D600F71900772A00772500776A00772500776D006F620077560077A20077560077A20077560076D600F75100772A00775100772A00775100772D006F6100770200774A0077100077240077A100770200774D006EE80077A100776A0077320074A0F08111241868C141101344218B60D186A99907F202A44207E865AFE01911E60128EBFFE8985608866B21089F8F1BFB8ED1D1270080B340C000004000000C1010CA98E36600000080000000400000000000E4062834061D10100802C1000000012000000321E981419114C908C092647C604435245500285500EE5538002055120230065407004A020288C45802008A260008220A882010882200B062008822E1890C0C0D11B4B7804000000B1180000A50000003308801CC4E11C6614013D88433884C38C4280077978077398710CE6000FED100EF4800E330C421EC2C11DCEA11C6630053D88433884831BCC033DC8433D8C033DCC788C7470077B08077948877070077A700376788770208719CC110EEC900EE1300F6E300FE3F00EF0500E3310C41DDE211CD8211DC2611E6630893BBC833BD04339B4033CBC833C84033BCCF0147660077B680737688772680737808770908770600776280776F8057678877780875F08877118877298877998812CEEF00EEEE00EF5C00EEC300362C8A11CE4A11CCCA11CE4A11CDC611CCA211CC4811DCA6106D6904339C84339984339C84339B8C33894433888033B94C32FBC833CFC823BD4033BB0C30CC7698770588772708374680778608774188774A08719CE530FEE000FF2500EE4900EE3400FE1200EEC500E3320281DDCC11EC2411ED2211CDC811EDCE01CE4E11DEA011E66185138B0433A9C833BCC50247660077B68073760877778077898514CF4900FF0500E331E6A1ECA611CE8211DDEC11D7E011EE4A11CCC211DF0610654858338CCC33BB0433DD04339FCC23CE4433B88C33BB0C38CC50A877998877718877408077A28077298815CE3100EECC00EE5500EF33023C1D2411EE4E117D8E11DDE011E6648193BB0833DB4831B84C3388C4339CCC33CB8C139C8C33BD4033CCC48B471080776600771088771588719DBC60EEC600FEDE006F0200FE5300FE5200FF6500E6E100EE3300EE5300FF3E006E9E00EE4500EF83023E2EC611CC2811DD8E117EC211DE6211DC4211DD8211DE8211F66209D3BBC433DB80339948339CC58BC7070077778077A08077A488777708719CE870EE5100EF0100EECC00EEF300EF3900EF4500E33283008877490073730877A708771A08774780777F88573908777A807789807000000007920000023010000721E482043880C19097232482023818C9191D144A01028643C3132428E9021A3580660B0570000008BD258D8066D501C141B4706D3232111254895E2144374458C7445625050000053444B2056657273696F6E77636861725F73697A656672616D652D706F696E7465726169722E6D61785F6465766963655F627566666572736169722E6D61785F636F6E7374616E745F627566666572736169722E6D61785F74687265616467726F75705F627566666572736169722E6D61785F74657874757265736169722E6D61785F726561645F77726974655F74657874757265736169722E6D61785F73616D706C6572734170706C65206D6574616C2076657273696F6E2033323032332E34303420286D6574616C66652D33323032332E343034294D6574616C6169722E636F6D70696C652E64656E6F726D735F64697361626C656169722E636F6D70696C652E666173745F6D6174685F656E61626C656169722E636F6D70696C652E6672616D656275666665725F66657463685F656E61626C656169722E696E6469726563745F6275666665726169722E6275666665725F73697A656169722E6C6F636174696F6E5F696E6465786169722E726561646169722E616464726573735F73706163656169722E7374727563745F747970655F696E666F5F365F7265736F757263655F7661725F305F305F6169722E696E6469726563745F617267756D656E746169722E6275666665726169722E726561645F7772697465666C6F61745F6D306169722E6172675F747970655F73697A656169722E6172675F747970655F616C69676E5F73697A656169722E6172675F747970655F6E616D656169722E6172675F6E616D655F7265736F757263655F7661725F305F315F73707644657363726970746F725365744275666665723073707644657363726970746F72536574306169722E7468726561645F706F736974696F6E5F696E5F74687265616467726F757075696E7433676C5F4C6F63616C496E766F636174696F6E494400C6750000000000003082A00823080F3482A00C23080A3182A01423088A3182A01C23080A3282A02423088A3282A02C23084D30C36006C119CC30A0819006330C6A3098C10C831A106B30C3A006C51ACC30A881B106330C6A70B0C10C831A206D30C3A006891BCC1028330C66A0066F3003B1C08119BCC10C013343D0CC1038330033146F900668A06D33486F80456F8006D9840694287069D0A581577DD60C121A60111AA04136A101250A5C1A7469E0551F18CCE0BC411BBC41655DA3D0066DF00615185CA43003F5060FE406D11BA081349901550A9C1B746DE085C1270633186830061E197C653083600AA730C31007A1800A470600C7711CC7711CC7719C1BB8811BB8811BB8811BB8811B5874A007966541A060166CA113A4C10A6820238109CA888DCDAECDA5ED8DAC8EADCCC58C2DEC6C6E94240EE4600EE8A00EECE00EF0200F526163B36B7349232B73A31B25D0835CC2D2E45CECCAE4E6D2DEDC4609F620A9B0343917B630B7B3BAB0B3B22FBB32B9B9B437B751023EC8292C4DCE65ECAD0D2E8DADECEB0D8E2EEDCD6D6E94A10FFCE00F52094B9373592B93732B631B254005000000A9180000250000000B0A7228877780077A587098433DB8C338B04339D0C382E61CC6A10DE8411EC2C11DE6211DE8211DDEC11D1634E3600EE7500FE1200FE4400FE1200FE7500EF4B08081077928877060077678877108077A28077258709CC338B4013BA4833D94C3026B1CD8211CDCE11CDC201CE4611CDC201CE8811EC2611CD0A11CC8611CC2811DD861C1010FF4200FE1500FF4800E00000000D11000000600000007CC3CA4833B9C033B94033DA0833C94433890C301000000612000006E0000001304442C1000000002000000A40AA87200000000F1300000470000002247C890511E485E00000000DDF23C7F605900005F5A5453323373707644657363726970746F7253657442756666657230616E7920706F696E7465726F6D6E69706F74656E74206368617253696D706C6520432B2B20544241416169722D616C6961732D73636F706573286D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F663332296169722D616C6961732D73636F70652D6172672830296C6C766D2E6C6F6F702E6D75737470726F6772657373000000130488992040CD0A61165C61C5400BB3F00A2B845AF8850D412B6C1858C1166261C3B00AB7100B1B0A55C08558C00559D830E4022EC4C2868016360CB9800BB2B061D06E2116360CBDD00BB1B02180050000000063084332DC10486830CB1008414783B134141CD2D18034160587CC120C3A50A13041346400A31BA288A6800230BA20BEA91A2800038341FCBB21100B33860D88401800442010186E482034986518842020060000090000005B0A22C8854317B61446B00B872E6C1998A017B60C4ED00B5B8628E8852D0325D1020000000000007120000003000000320E10228401B2060000000000000000650C00002500000012039428010000000300000034000000090000004C000000010000005800000000000000580000000200000088000000000000003D00000018000000070000002D00000007000000000000008800000000000000000000000200000000000000070000002D000000070000002D000000FFFFFFFF0024000000000000070000000000000007000000FFFFFFFF00100000000000005D0C000019000000120394B5000000005F5A4C335F31396D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F66333233323032332E34303461697236342D6170706C652D6D61636F737831352E302E3000000000000000000000000048F9FFFF0000030004000000A606000023707261676D6120636C616E6720646961676E6F737469632069676E6F72656420222D576D697373696E672D70726F746F7479706573220A23707261676D6120636C616E6720646961676E6F737469632069676E6F72656420222D576D697373696E672D627261636573220A0A23696E636C756465203C6D6574616C5F7374646C69623E0A23696E636C756465203C73696D642F73696D642E683E0A0A7573696E67206E616D657370616365206D6574616C3B0A0A74656D706C6174653C747970656E616D6520542C2073697A655F74204E756D3E0A73747275637420737076556E7361666541727261790A7B0A202020205420656C656D656E74735B4E756D203F204E756D203A20315D3B0A202020200A20202020746872656164205426206F70657261746F72205B5D202873697A655F7420706F7329207468726561640A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A20202020636F6E73746578707220636F6E737420746872656164205426206F70657261746F72205B5D202873697A655F7420706F732920636F6E7374207468726561640A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A202020200A20202020646576696365205426206F70657261746F72205B5D202873697A655F7420706F7329206465766963650A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A20202020636F6E73746578707220636F6E737420646576696365205426206F70657261746F72205B5D202873697A655F7420706F732920636F6E7374206465766963650A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A202020200A20202020636F6E73746578707220636F6E737420636F6E7374616E74205426206F70657261746F72205B5D202873697A655F7420706F732920636F6E737420636F6E7374616E740A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A202020200A2020202074687265616467726F7570205426206F70657261746F72205B5D202873697A655F7420706F73292074687265616467726F75700A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A20202020636F6E73746578707220636F6E73742074687265616467726F7570205426206F70657261746F72205B5D202873697A655F7420706F732920636F6E73742074687265616467726F75700A202020207B0A202020202020202072657475726E20656C656D656E74735B706F735D3B0A202020207D0A7D3B0A0A737472756374205F360A7B0A20202020666C6F6174205F6D305B315D3B0A7D3B0A0A7374727563742073707644657363726970746F72536574427566666572300A7B0A20202020646576696365205F362A205F7265736F757263655F7661725F305F305F205B5B69642830295D5D3B0A20202020646576696365205F362A205F7265736F757263655F7661725F305F315F205B5B69642831295D5D3B0A7D3B0A0A636F6E7374616E7420737076556E7361666541727261793C666C6F61742C20343E205F3139203D20737076556E7361666541727261793C666C6F61742C20343E287B20342E302C20352E302C20362E302C20372E30207D293B0A0A6B65726E656C20766F6964206D756C74695F6465766963655F6D756C5F64697370617463685F315F656C656D656E74776973655F345F66333228636F6E7374616E742073707644657363726970746F7253657442756666657230262073707644657363726970746F7253657430205B5B6275666665722830295D5D2C2075696E743320676C5F4C6F63616C496E766F636174696F6E4944205B5B7468726561645F706F736974696F6E5F696E5F74687265616467726F75705D5D290A7B0A20202020666F72202875696E74205F3239203D20676C5F4C6F63616C496E766F636174696F6E49442E783B20696E74285F323929203C20696E74283475293B20290A202020207B0A2020202020202020282A73707644657363726970746F72536574302E5F7265736F757263655F7661725F305F315F292E5F6D305B5F32395D203D20282A73707644657363726970746F72536574302E5F7265736F757263655F7661725F305F305F292E5F6D305B5F32395D202A205F31395B5F32395D3B0A20202020202020205F3239202B3D203332753B0A2020202020202020636F6E74696E75653B0A202020207D0A7D0A0A000008000C000400080014001C000000040000000800000000001400180000000000"> : vector<6336xi8>
vm.global.ref private mutable @__device_a_executable_0_multi_device_mul_dispatch_0 : !vm.ref<!hal.executable>
vm.global.ref private mutable @device_b : !vm.ref<!hal.device>
vm.global.ref private mutable @__device_b_executable_0_multi_device_mul_dispatch_1 : !vm.ref<!hal.executable>
vm.func private @__multi_device_mul_memoize_apply() -> !vm.ref<!hal.command_buffer> attributes {inlining_policy = #util.inline.never} {
%c13 = vm.const.i32 13
%c28 = vm.const.i32 28
%c2 = vm.const.i32 2
%null = vm.const.ref.zero : !vm.ref<!hal.buffer>
%c1 = vm.const.i32 1
%c3 = vm.const.i32 3
%zero = vm.const.i32.zero
%c64 = vm.const.i64 64
%c16 = vm.const.i64 16
%zero_0 = vm.const.i64.zero
%c-1 = vm.const.i64 -1
%device_a = vm.global.load.ref @device_a : !vm.ref<!hal.device>
%__device_a_executable_0_multi_device_mul_dispatch_0 = vm.global.load.ref @__device_a_executable_0_multi_device_mul_dispatch_0 : !vm.ref<!hal.executable>
%ref = vm.call @hal.command_buffer.create(%device_a, %zero, %c3, %c-1, %c3) : (!vm.ref<!hal.device>, i32, i32, i64, i32) -> !vm.ref<!hal.command_buffer>
vm.call.variadic @hal.command_buffer.dispatch(%ref, %__device_a_executable_0_multi_device_mul_dispatch_0, %zero, %c1, %c1, %c1, %zero_0, [], [(%zero, %zero, %null, %zero_0, %c16), (%zero, %c2, %null, %zero_0, %c64)]) : (!vm.ref<!hal.command_buffer>, !vm.ref<!hal.executable>, i32, i32, i32, i32, i64, i32 ..., tuple<i32, i32, !vm.ref<!hal.buffer>, i64, i64> ...)
vm.call @hal.command_buffer.execution_barrier(%ref, %c28, %c13, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, i64) -> ()
vm.call @hal.command_buffer.copy_buffer(%ref, %c2, %c1, %null, %zero_0, %null, %zero_0, %c16, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, !vm.ref<!hal.buffer>, i64, !vm.ref<!hal.buffer>, i64, i64, i64) -> ()
vm.call @hal.command_buffer.execution_barrier(%ref, %c28, %c13, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, i64) -> ()
vm.call @hal.command_buffer.finalize(%ref) : (!vm.ref<!hal.command_buffer>) -> ()
vm.return %ref : !vm.ref<!hal.command_buffer>
}
vm.global.ref private mutable @__multi_device_mul_memoize_result_0_device_a : !vm.ref<!hal.command_buffer>
vm.func private @__multi_device_mul_memoize_apply_0() -> !vm.ref<!hal.command_buffer> attributes {inlining_policy = #util.inline.never} {
%c13 = vm.const.i32 13
%c28 = vm.const.i32 28
%c2 = vm.const.i32 2
%null = vm.const.ref.zero : !vm.ref<!hal.buffer>
%c1 = vm.const.i32 1
%c3 = vm.const.i32 3
%zero = vm.const.i32.zero
%c64 = vm.const.i64 64
%c16 = vm.const.i64 16
%zero_0 = vm.const.i64.zero
%c-1 = vm.const.i64 -1
%device_b = vm.global.load.ref @device_b : !vm.ref<!hal.device>
%__device_b_executable_0_multi_device_mul_dispatch_1 = vm.global.load.ref @__device_b_executable_0_multi_device_mul_dispatch_1 : !vm.ref<!hal.executable>
%ref = vm.call @hal.command_buffer.create(%device_b, %zero, %c3, %c-1, %c3) : (!vm.ref<!hal.device>, i32, i32, i64, i32) -> !vm.ref<!hal.command_buffer>
vm.call.variadic @hal.command_buffer.dispatch(%ref, %__device_b_executable_0_multi_device_mul_dispatch_1, %zero, %c1, %c1, %c1, %zero_0, [], [(%zero, %zero, %null, %zero_0, %c16), (%zero, %c2, %null, %zero_0, %c64)]) : (!vm.ref<!hal.command_buffer>, !vm.ref<!hal.executable>, i32, i32, i32, i32, i64, i32 ..., tuple<i32, i32, !vm.ref<!hal.buffer>, i64, i64> ...)
vm.call @hal.command_buffer.execution_barrier(%ref, %c28, %c13, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, i64) -> ()
vm.call @hal.command_buffer.copy_buffer(%ref, %c2, %c1, %null, %zero_0, %null, %zero_0, %c16, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, !vm.ref<!hal.buffer>, i64, !vm.ref<!hal.buffer>, i64, i64, i64) -> ()
vm.call @hal.command_buffer.execution_barrier(%ref, %c28, %c13, %zero_0) : (!vm.ref<!hal.command_buffer>, i32, i32, i64) -> ()
vm.call @hal.command_buffer.finalize(%ref) : (!vm.ref<!hal.command_buffer>) -> ()
vm.return %ref : !vm.ref<!hal.command_buffer>
}
vm.global.ref private mutable @__multi_device_mul_memoize_result_0_device_b : !vm.ref<!hal.command_buffer>
vm.import private @hal.buffer.assert(%buffer : !vm.ref<!hal.buffer>, %message : !vm.buffer, %allocator : !vm.ref<!hal.allocator>, %minimum_length : i64, %memory_types : i32, %buffer_usage : i32)
vm.import private @hal.buffer_view.create(%buffer : !vm.ref<!hal.buffer>, %source_offset : i64, %source_length : i64, %element_type : i32, %encoding_type : i32, %shape : i64 ...) -> !vm.ref<!hal.buffer_view> attributes {nosideeffects}
vm.import private @hal.buffer_view.assert(%buffer_view : !vm.ref<!hal.buffer_view>, %message : !vm.buffer, %element_type : i32, %encoding_type : i32, %shape : i64 ...)
vm.import private @hal.buffer_view.buffer(%buffer_view : !vm.ref<!hal.buffer_view>) -> !vm.ref<!hal.buffer> attributes {nosideeffects}
vm.import private @hal.command_buffer.create(%device : !vm.ref<!hal.device>, %modes : i32, %command_categories : i32, %queue_affinity : i64, %binding_capacity : i32) -> !vm.ref<!hal.command_buffer> attributes {minimum_version = 6 : i32}
vm.import private @hal.command_buffer.finalize(%command_buffer : !vm.ref<!hal.command_buffer>)
vm.import private @hal.command_buffer.execution_barrier(%command_buffer : !vm.ref<!hal.command_buffer>, %source_stage_mask : i32, %target_stage_mask : i32, %flags : i64)
vm.import private @hal.command_buffer.copy_buffer(%command_buffer : !vm.ref<!hal.command_buffer>, %source_buffer_slot : i32, %target_buffer_slot : i32, %source_buffer : !vm.ref<!hal.buffer>, %source_offset : i64, %target_buffer : !vm.ref<!hal.buffer>, %target_offset : i64, %length : i64, %flags : i64)
vm.import private @hal.command_buffer.dispatch(%command_buffer : !vm.ref<!hal.command_buffer>, %executable : !vm.ref<!hal.executable>, %entry_point : i32, %workgroup_x : i32, %workgroup_y : i32, %workgroup_z : i32, %flags : i64, %constants : i32 ..., %bindings : tuple<i32, i32, !vm.ref<!hal.buffer>, i64, i64> ...)
vm.import private @hal.device.allocator(%device : !vm.ref<!hal.device>) -> !vm.ref<!hal.allocator> attributes {nosideeffects}
vm.import private @hal.device.query.i64(%device : !vm.ref<!hal.device>, %category : !vm.buffer, %key : !vm.buffer) -> (i32, i64) attributes {nosideeffects}
vm.import private @hal.device.queue.alloca(%device : !vm.ref<!hal.device>, %queue_affinity : i64, %wait_fence : !vm.ref<!hal.fence>, %signal_fence : !vm.ref<!hal.fence>, %pool : i64, %memory_types : i32, %buffer_usage : i32, %allocation_size : i64, %flags : i64) -> !vm.ref<!hal.buffer>
vm.import private @hal.device.queue.dealloca(%device : !vm.ref<!hal.device>, %queue_affinity : i64, %wait_fence : !vm.ref<!hal.fence>, %signal_fence : !vm.ref<!hal.fence>, %buffer : !vm.ref<!hal.buffer>, %flags : i64)
vm.import private @hal.device.queue.execute.indirect(%device : !vm.ref<!hal.device>, %queue_affinity : i64, %wait_fence : !vm.ref<!hal.fence>, %signal_fence : !vm.ref<!hal.fence>, %command_buffer : !vm.ref<!hal.command_buffer>, %flags : i64, %binding_table : tuple<!vm.ref<!hal.buffer>, i64, i64> ...)
vm.import private @hal.devices.count() -> i32 attributes {nosideeffects}
vm.import private @hal.devices.get(%index : i32) -> !vm.ref<!hal.device> attributes {nosideeffects}
vm.import private @hal.executable.create(%device : !vm.ref<!hal.device>, %queue_affinity : i64, %executable_format : !vm.buffer, %executable_data : !vm.buffer, %constants : !vm.buffer) -> !vm.ref<!hal.executable> attributes {nosideeffects}
vm.import private @hal.fence.create(%device : !vm.ref<!hal.device>, %flags : i64) -> !vm.ref<!hal.fence>
vm.import private @hal.fence.join(%flags : i64, %fences : !vm.ref<!hal.fence> ...) -> !vm.ref<!hal.fence> attributes {nosideeffects}
vm.import private @hal.fence.signal(%fence : !vm.ref<!hal.fence>) -> () attributes {nosideeffects}
vm.import private @hal.fence.await(%timeout_millis : i32, %flags : i64, %fences : !vm.ref<!hal.fence> ...) -> i32 attributes {vm.yield}
vm.rodata private @_utf8_input0_2D1C6C95BF4E8E4E {alignment = 1 : i64} "input0"
vm.rodata private @_utf8_tensor_8F58A528F0C2C1FD {alignment = 1 : i64} "tensor"
vm.func private @multi_device_mul(%arg0: !vm.ref<!hal.buffer_view>, %arg1: !vm.ref<!hal.fence>, %arg2: !vm.ref<!hal.fence>) -> !vm.ref<!hal.buffer_view> attributes {iree.reflection = {iree.abi.declaration = "async func @multi_device_mul(%input0: tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>}) -> (%output0: tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>})", iree.abi.model = "coarse-fences"}} {
%c48 = vm.const.i32 48
%c3075 = vm.const.i32 3075
%c16 = vm.const.i32 16
%c1 = vm.const.i32 1
%c553648160 = vm.const.i32 553648160
%c4 = vm.const.i64 4
%c16_0 = vm.const.i64 16
%zero = vm.const.i64.zero
%c64 = vm.const.i64 64
%c-1 = vm.const.i64 -1
%device_a = vm.global.load.ref @device_a : !vm.ref<!hal.device>
%device_b = vm.global.load.ref @device_b : !vm.ref<!hal.device>
%__multi_device_mul_memoize_result_0_device_a = vm.global.load.ref @__multi_device_mul_memoize_result_0_device_a : !vm.ref<!hal.command_buffer>
%__multi_device_mul_memoize_result_0_device_b = vm.global.load.ref @__multi_device_mul_memoize_result_0_device_b : !vm.ref<!hal.command_buffer>
%_utf8_input0_2D1C6C95BF4E8E4E = vm.const.ref.rodata @_utf8_input0_2D1C6C95BF4E8E4E : !vm.buffer
vm.call.variadic @hal.buffer_view.assert(%arg0, %_utf8_input0_2D1C6C95BF4E8E4E, %c553648160, %c1, [%c4]) : (!vm.ref<!hal.buffer_view>, !vm.buffer, i32, i32, i64 ...)
%ref = vm.call @hal.buffer_view.buffer(%arg0) {nosideeffects} : (!vm.ref<!hal.buffer_view>) -> !vm.ref<!hal.buffer>
%ref_1 = vm.call @hal.device.allocator(%device_a) {nosideeffects} : (!vm.ref<!hal.device>) -> !vm.ref<!hal.allocator>
%_utf8_tensor_8F58A528F0C2C1FD = vm.const.ref.rodata @_utf8_tensor_8F58A528F0C2C1FD : !vm.buffer
vm.call @hal.buffer.assert(%ref, %_utf8_tensor_8F58A528F0C2C1FD, %ref_1, %c16_0, %c16, %c3075) : (!vm.ref<!hal.buffer>, !vm.buffer, !vm.ref<!hal.allocator>, i64, i32, i32) -> ()
%ref_2 = vm.call @hal.fence.create(%device_b, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
%ref_3 = vm.call @hal.device.queue.alloca(%device_b, %c-1, %arg1, %ref_2, %zero, %c48, %c3075, %c16_0, %zero) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, i64, i32, i32, i64, i64) -> !vm.ref<!hal.buffer>
%join_b0 = vm.call.variadic @hal.fence.join(%zero, [%ref_2]) {nosideeffects} : (i64, !vm.ref<!hal.fence> ...) -> !vm.ref<!hal.fence>
%ref_4 = vm.call @hal.fence.create(%device_a, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
%ref_5 = vm.call @hal.device.queue.alloca(%device_a, %c-1, %arg1, %ref_4, %zero, %c48, %c3075, %c64, %zero) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, i64, i32, i32, i64, i64) -> !vm.ref<!hal.buffer>
%ref_6 = vm.call.variadic @hal.fence.join(%zero, [%ref_4]) {nosideeffects} : (i64, !vm.ref<!hal.fence> ...) -> !vm.ref<!hal.fence>
%ref_7 = vm.call @hal.fence.create(%device_a, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
vm.call.variadic @hal.device.queue.execute.indirect(%device_a, %c-1, %ref_6, %ref_7, %__multi_device_mul_memoize_result_0_device_a, %zero, [(%ref, %zero, %c16_0), (%ref_3, %zero, %c16_0), (%ref_5, %zero, %c64)]) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, !vm.ref<!hal.command_buffer>, i64, tuple<!vm.ref<!hal.buffer>, i64, i64> ...)
%ref_8 = vm.call @hal.fence.create(%device_b, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
vm.call @hal.device.queue.dealloca(%device_a, %c-1, %ref_7, %ref_8, %ref_5, %zero) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, !vm.ref<!hal.buffer>, i64) -> ()
%ref_9 = vm.call @hal.fence.create(%device_a, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
%ref_10 = vm.call @hal.device.queue.alloca(%device_a, %c-1, %ref_8, %ref_9, %zero, %c48, %c3075, %c16_0, %zero) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, i64, i32, i32, i64, i64) -> !vm.ref<!hal.buffer>
%join_b1 = vm.call.variadic @hal.fence.join(%zero, [%ref_9]) {nosideeffects} : (i64, !vm.ref<!hal.fence> ...) -> !vm.ref<!hal.fence>
%ref_11 = vm.call @hal.fence.create(%device_b, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
%ref_12 = vm.call @hal.device.queue.alloca(%device_b, %c-1, %ref_8, %ref_11, %zero, %c48, %c3075, %c64, %zero) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, i64, i32, i32, i64, i64) -> !vm.ref<!hal.buffer>
%ref_13 = vm.call.variadic @hal.fence.join(%zero, [ %ref_11]) {nosideeffects} : (i64, !vm.ref<!hal.fence> ...) -> !vm.ref<!hal.fence>
%ref_14 = vm.call @hal.fence.create(%device_b, %zero) : (!vm.ref<!hal.device>, i64) -> !vm.ref<!hal.fence>
vm.call.variadic @hal.device.queue.execute.indirect(%device_b, %c-1, %ref_13, %ref_14, %__multi_device_mul_memoize_result_0_device_b, %zero, [(%ref_3, %zero, %c16_0), (%ref_10, %zero, %c16_0), (%ref_12, %zero, %c64)]) : (!vm.ref<!hal.device>, i64, !vm.ref<!hal.fence>, !vm.ref<!hal.fence>, !vm.ref<!hal.command_buffer>, i64, tuple<!vm.ref<!hal.buffer>, i64, i64> ...)
%0 = vm.call.variadic @hal.fence.await(%c-1, %zero, [%ref_14]) : (i64, i64, !vm.ref<!hal.fence> ...) -> i32
vm.call @hal.fence.signal(%arg2) : (!vm.ref<!hal.fence>) -> ()
%ref_15 = vm.call.variadic @hal.buffer_view.create(%ref_10, %zero, %c16_0, %c553648160, %c1, [%c4]) {nosideeffects} : (!vm.ref<!hal.buffer>, i64, i64, i32, i32, i64 ...) -> !vm.ref<!hal.buffer_view>
vm.return %ref_15 : !vm.ref<!hal.buffer_view>
}
vm.export @multi_device_mul attributes {iree.abi.stub, iree.reflection = {iree.abi.declaration = "async func @multi_device_mul(%input0: tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>}) -> (%output0: tensor<4xf32> {iree.abi.affinity = #hal.device.promise<@device_a>})", iree.abi.model = "coarse-fences"}}
vm.export @__init
vm.func private @__init() {
%c1 = vm.const.i32 1
%null = vm.const.ref.zero : !vm.buffer
%c14 = vm.const.i32 14
%c-1 = vm.const.i64 -1
%c18 = vm.const.i32 18
%zero = vm.const.i32.zero
%zero_0 = vm.const.i64.zero
%c1_1 = vm.const.i64 1
%null_2 = vm.const.ref.zero : !vm.ref<!hal.device>
%0 = vm.call @hal.devices.count() {nosideeffects} : () -> i32
%1 = vm.ext.i32.i64.s %0 : i32 -> i64
vm.br ^bb1(%zero_0, %zero_0, %null_2 : i64, i64, !vm.ref<!hal.device>)
^bb1(%2: i64, %3: i64, %4: !vm.ref<!hal.device>): // 2 preds: ^bb0, ^bb4
%rnz = vm.cmp.nz.ref %4 : !vm.ref<!hal.device>
%5 = vm.xor.i32 %rnz, %c1 : i32
%slt = vm.cmp.lt.i64.s %2, %1 : i64
%6 = vm.and.i32 %5, %slt : i32
vm.cond_br %6, ^bb2, ^bb5
^bb2: // pred: ^bb1
%7 = vm.trunc.i64.i32 %2 : i64 -> i32
%ref = vm.call @hal.devices.get(%7) {nosideeffects} : (i32) -> !vm.ref<!hal.device>
%_utf8_hal_device_id_369FCE3B885986F6 = vm.const.ref.rodata @_utf8_hal_device_id_369FCE3B885986F6 : !vm.buffer
%_utf8_local_2C7344D4E05782F4 = vm.const.ref.rodata @_utf8_local_2C7344D4E05782F4 : !vm.buffer
%8:2 = vm.call @hal.device.query.i64(%ref, %_utf8_hal_device_id_369FCE3B885986F6, %_utf8_local_2C7344D4E05782F4) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz = vm.cmp.nz.i64 %8#1 : i64
%9 = vm.select.i32 %8#0, %nz, %zero : i32
vm.cond_br %9, ^bb3, ^bb4(%zero : i32)
^bb3: // pred: ^bb2
%_utf8_hal_executable_format_C7128E2AE1BE720D = vm.const.ref.rodata @_utf8_hal_executable_format_C7128E2AE1BE720D : !vm.buffer
%_utf8_vmvx_bytecode_fb_90AB3AF67FE09641 = vm.const.ref.rodata @_utf8_vmvx_bytecode_fb_90AB3AF67FE09641 : !vm.buffer
%10:2 = vm.call @hal.device.query.i64(%ref, %_utf8_hal_executable_format_C7128E2AE1BE720D, %_utf8_vmvx_bytecode_fb_90AB3AF67FE09641) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz_3 = vm.cmp.nz.i64 %10#1 : i64
%11 = vm.select.i32 %10#0, %nz_3, %zero : i32
vm.br ^bb4(%11 : i32)
^bb4(%12: i32): // 2 preds: ^bb2, ^bb3
%eq = vm.cmp.eq.i64 %3, %zero_0 : i64
%13 = vm.select.i64 %12, %c1_1, %zero_0 : i64
%14 = vm.add.i64 %3, %13 : i64
%15 = vm.and.i32 %12, %eq : i32
%ref_4 = vm.select.ref %15, %ref, %null_2 : !vm.ref<!hal.device>
%16 = vm.add.i64 %2, %c1_1 : i64
vm.br ^bb1(%16, %14, %ref_4 : i64, i64, !vm.ref<!hal.device>)
^bb5: // pred: ^bb1
vm.cond_br %5, ^bb6, ^bb7
^bb6: // pred: ^bb5
vm.fail %c18, "HAL device `device_a` not found or unavailable: #hal.device.target<\22local\22, {ordinal = 0 : index}, [#hal.executable.target<\22vmvx\22, \22vmvx-bytecode-fb\22, {iree.encoding.resolver = #iree_cpu.vmvx_encoding_layout<>, ukernels = \22none\22}>]>"
^bb7: // pred: ^bb5
%_utf8_hal_executable_format_C7128E2AE1BE720D_5 = vm.const.ref.rodata @_utf8_hal_executable_format_C7128E2AE1BE720D : !vm.buffer
%_utf8_vmvx_bytecode_fb_90AB3AF67FE09641_6 = vm.const.ref.rodata @_utf8_vmvx_bytecode_fb_90AB3AF67FE09641 : !vm.buffer
%17:2 = vm.call @hal.device.query.i64(%4, %_utf8_hal_executable_format_C7128E2AE1BE720D_5, %_utf8_vmvx_bytecode_fb_90AB3AF67FE09641_6) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz_7 = vm.cmp.nz.i64 %17#1 : i64
%18 = vm.select.i32 %17#0, %nz_7, %zero : i32
%19 = vm.select.i64 %18, %zero_0, %c-1 : i64
%eq_8 = vm.cmp.eq.i64 %19, %zero_0 : i64
vm.global.store.ref %4, @device_a : !vm.ref<!hal.device>
vm.cond_br %eq_8, ^bb8, ^bb9
^bb8: // pred: ^bb7
%multi_device_mul_dispatch_0_vmvx_bytecode_fb = vm.const.ref.rodata @multi_device_mul_dispatch_0_vmvx_bytecode_fb : !vm.buffer
%ref_9 = vm.call @hal.executable.create(%4, %c-1, %_utf8_vmvx_bytecode_fb_90AB3AF67FE09641_6, %multi_device_mul_dispatch_0_vmvx_bytecode_fb, %null) {nosideeffects} : (!vm.ref<!hal.device>, i64, !vm.buffer, !vm.buffer, !vm.buffer) -> !vm.ref<!hal.executable>
vm.global.store.ref %ref_9, @__device_a_executable_0_multi_device_mul_dispatch_0 : !vm.ref<!hal.executable>
vm.br ^bb10(%zero_0, %zero_0, %null_2 : i64, i64, !vm.ref<!hal.device>)
^bb9: // pred: ^bb7
vm.fail %c14, "HAL device `device_a` does not support any variant of executable `multi_device_mul_dispatch_0`; available formats: [vmvx-bytecode-fb]"
^bb10(%20: i64, %21: i64, %22: !vm.ref<!hal.device>): // 2 preds: ^bb8, ^bb13
%rnz_10 = vm.cmp.nz.ref %22 : !vm.ref<!hal.device>
%23 = vm.xor.i32 %rnz_10, %c1 : i32
%slt_11 = vm.cmp.lt.i64.s %20, %1 : i64
%24 = vm.and.i32 %23, %slt_11 : i32
vm.cond_br %24, ^bb11, ^bb14
^bb11: // pred: ^bb10
%25 = vm.trunc.i64.i32 %20 : i64 -> i32
%ref_12 = vm.call @hal.devices.get(%25) {nosideeffects} : (i32) -> !vm.ref<!hal.device>
%_utf8_hal_device_id_369FCE3B885986F6_13 = vm.const.ref.rodata @_utf8_hal_device_id_369FCE3B885986F6 : !vm.buffer
%_utf8_metal_FB846B13C56E663B = vm.const.ref.rodata @_utf8_metal_FB846B13C56E663B : !vm.buffer
%26:2 = vm.call @hal.device.query.i64(%ref_12, %_utf8_hal_device_id_369FCE3B885986F6_13, %_utf8_metal_FB846B13C56E663B) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz_14 = vm.cmp.nz.i64 %26#1 : i64
%27 = vm.select.i32 %26#0, %nz_14, %zero : i32
vm.cond_br %27, ^bb12, ^bb13(%zero : i32)
^bb12: // pred: ^bb11
%_utf8_metal_msl_fb_BD2413475D7AB9A0 = vm.const.ref.rodata @_utf8_metal_msl_fb_BD2413475D7AB9A0 : !vm.buffer
%28:2 = vm.call @hal.device.query.i64(%ref_12, %_utf8_hal_executable_format_C7128E2AE1BE720D_5, %_utf8_metal_msl_fb_BD2413475D7AB9A0) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz_15 = vm.cmp.nz.i64 %28#1 : i64
%29 = vm.select.i32 %28#0, %nz_15, %zero : i32
vm.br ^bb13(%29 : i32)
^bb13(%30: i32): // 2 preds: ^bb11, ^bb12
%eq_16 = vm.cmp.eq.i64 %21, %zero_0 : i64
%31 = vm.select.i64 %30, %c1_1, %zero_0 : i64
%32 = vm.add.i64 %21, %31 : i64
%33 = vm.and.i32 %30, %eq_16 : i32
%ref_17 = vm.select.ref %33, %ref_12, %null_2 : !vm.ref<!hal.device>
%34 = vm.add.i64 %20, %c1_1 : i64
vm.br ^bb10(%34, %32, %ref_17 : i64, i64, !vm.ref<!hal.device>)
^bb14: // pred: ^bb10
vm.cond_br %23, ^bb15, ^bb16
^bb15: // pred: ^bb14
vm.fail %c18, "HAL device `device_b` not found or unavailable: #hal.device.target<\22metal\22, [#hal.executable.target<\22metal-spirv\22, \22metal-msl-fb\22, {iree.gpu.target = #iree_gpu.target<arch = \22apple\22, features = \22spirv:v1.3,cap:Shader\22, wgp = <compute = fp32|fp16|int64|int32|int16|int8, storage = b64|b32|b16|b8, subgroup = shuffle|arithmetic, dot = dp4xi8toi32, mma = [], subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 1024], max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 32768, max_workgroup_counts = [65535, 65535, 65535]>>}>]>"
^bb16: // pred: ^bb14
%_utf8_metal_msl_fb_BD2413475D7AB9A0_18 = vm.const.ref.rodata @_utf8_metal_msl_fb_BD2413475D7AB9A0 : !vm.buffer
%35:2 = vm.call @hal.device.query.i64(%22, %_utf8_hal_executable_format_C7128E2AE1BE720D_5, %_utf8_metal_msl_fb_BD2413475D7AB9A0_18) {nosideeffects} : (!vm.ref<!hal.device>, !vm.buffer, !vm.buffer) -> (i32, i64)
%nz_19 = vm.cmp.nz.i64 %35#1 : i64
%36 = vm.select.i32 %35#0, %nz_19, %zero : i32
%37 = vm.select.i64 %36, %zero_0, %c-1 : i64
%eq_20 = vm.cmp.eq.i64 %37, %zero_0 : i64
vm.global.store.ref %22, @device_b : !vm.ref<!hal.device>
vm.cond_br %eq_20, ^bb17, ^bb18
^bb17: // pred: ^bb16
%multi_device_mul_dispatch_1_metal_msl_fb = vm.const.ref.rodata @multi_device_mul_dispatch_1_metal_msl_fb : !vm.buffer
%ref_21 = vm.call @hal.executable.create(%22, %c-1, %_utf8_metal_msl_fb_BD2413475D7AB9A0_18, %multi_device_mul_dispatch_1_metal_msl_fb, %null) {nosideeffects} : (!vm.ref<!hal.device>, i64, !vm.buffer, !vm.buffer, !vm.buffer) -> !vm.ref<!hal.executable>
vm.global.store.ref %ref_21, @__device_b_executable_0_multi_device_mul_dispatch_1 : !vm.ref<!hal.executable>
%ref_22 = vm.call @__multi_device_mul_memoize_apply() : () -> !vm.ref<!hal.command_buffer>
vm.global.store.ref %ref_22, @__multi_device_mul_memoize_result_0_device_a : !vm.ref<!hal.command_buffer>
%ref_23 = vm.call @__multi_device_mul_memoize_apply_0() : () -> !vm.ref<!hal.command_buffer>
vm.global.store.ref %ref_23, @__multi_device_mul_memoize_result_0_device_b : !vm.ref<!hal.command_buffer>
vm.return
^bb18: // pred: ^bb16
vm.fail %c14, "HAL device `device_b` does not support any variant of executable `multi_device_mul_dispatch_1`; available formats: [metal-msl-fb]"
}
}
}
Yep! The MVP only needs semaphore handling changes on the CPU side as it already uses the iree_hal_buffer_t vtable to perform the mapping instead of casting things directly. Import/export are required for non-CPU targets that want to share resources that have no host pointers - for example, Vulkan and CUDA/HIP. My expectation for that case is that either the import/export need to be extremely fast (implemented in drivers as a pointer indirection) or extremely infrequent (done as part of startup/loading tasks, etc). To have import/export work in mainline latency-critical execution it'll likely need a cache on the HAL side of imported buffers and that may require weak references so it's a bit of a ways off.
I’ve been looking into the semaphore timepoint export/import. We currently don’t use/have HIP available, but we’d be happy to contribute vulkan & local_* to the initial MVP. I’d like to start from the vulkan -> local_* angle, if that sounds good to you?
For vulkan, timepoint export appears to be “straight forward”, through vkGetSemaphoreFdKHR & iree_make_wait_primitive. (Could possibly use the timepoint semaphore itself to wait on for the cpu, but looks like a potential optimization to me)
Importing the timepoint into local_task, I found the following steps might work:
iree_hal_semaphore_acquire_timepointof thevalueiree_wait_handle_wrap_primitiveof thewait_primitiveto import.- Store the
wait_handleinto theiree_hal_task_timepoint_t.event
Does that sound correct? Anything else I am missing?
Thanks a lot for taking the time to outline this. We would be happy to help with any issues you created here.
We’ve been doing some prototyping ourselves using the metal/local-task drivers, where we took a shortcut for the “programs will be able to selectively allocate buffers that are compatible with all devices” by simply allocating everything as shared using the metal allocator and manually working around the semaphore issues by hand-modifying the HAL IR.
I have some questions about how to handle buffer import/export—or whether it's even necessary at all for the MVP. My observation so far is that drivers typically require buffers to be allocated using their own allocator, since many of the
iree_hal_*_command_buffer_*functions expect a native buffer handle. However, this doesn't seem to be the case for the local-task device, as I couldn’t find any references indicating that it requires its own native handle (which I assume isiree_hal_heap_buffer_t).Given this, is my understanding correct that for the MVP, buffer import/export might not be required at all? As you mentioned in (#20857/#20855), shared allocations would be performed by the device that supports them (metal, in this example), so the HAL functions on the GPU device would still have access to their native handles, and the local devices wouldn’t care?
If the import/export is still needed would it be something that is handled solely in the runtime or would the transfer ops have to be translated into some form of import/export ops instead of eliding them completely?
If its of any interest for someone, this is the IR snipped i got running for local-task/metal when replacing the local-task allocator with the metal allocator:
metal/local-task example originated from the
tools/test/iree-run-module-multi.mlir testcase, there are not a lot of modifications, mostly to resolve semaphore errors like joining on multiple fences that are not from the same device.compile command i used:
./build/tools/iree-compile --compile-from=vm --iree-execution-model=async-external --iree-hal-target-device=device_a=local[0] --iree-hal-target-device=device_b=metal --iree-hal-local-target-device-backends=vmvx working_multi_device.mlir -o multi.vmfbrun command:
./build/tools/iree-run-module --module=multi.vmfb --device=local-task --device=metal --input=4xf32=10,11,12,13 --function=multi_device_mul
Thanks for sharing this example. I'm trying to execute it in CPU+Metal with v3.4.0 but it's not working so far: I don't get the resulting hal.buffer_view. How did you set up your environment?
Importing the timepoint into
local_task, I found the following steps might work:
iree_hal_semaphore_acquire_timepointof thevalueiree_wait_handle_wrap_primitiveof thewait_primitiveto import.- Store the
wait_handleinto theiree_hal_task_timepoint_t.eventDoes that sound correct? Anything else I am missing?
Vulkan doesn't use the iree_hal_semaphore_acquire_timepoint API so that won't work there. Binary semaphores in Vulkan don't work either as they are auto-reset and don't allow multiple waiters (something we do). I've not looked in to how timeline semaphores are exported (via e.g. vkGetSemaphoreWin32HandleKHR) but that's where you'd want to start. If there's not a way to do that then it'll need a potentially complex implementation in the Vulkan HAL for managing native semaphores as well as exported timepoints in a way that does not induce host round-trips in the common case (and ideally not in the uncommon case). The entire Vulkan HAL needs a rewrite at some point, so adding a lot of new infra there is scary unless it can be quarantined and carried forward. But maybe there's a really simple solution and none of that's needed, just needs someone to do some experiments :)
Metal, CUDA, and HIP are much easier to start with because they do allow (though slow/higher latency) host callbacks to be inserted into the submission queue and those callbacks can then signal wait primitives directly.
@Achiirua if you want to execute it yourself it would require some local changes. I pushed them here: https://github.com/ziereis/iree/tree/metal_cpu_prototype. But keep in mind this is no proper solution to the general problem :D
with this i get this output:
ziereis@Thomass-MacBook-Pro ~/p/iree (metal_cpu_prototype)> ./build/tools/iree-run-module --module=multi.vmfb --device=local-task --device=metal --input=4xf32=10,11,12,13 --function=multi_device_mul --task_topology_group_count=1
Creating Metal HAL driver...
Metal driver created successfully.
Metal device created successfully.
Metal device created successfully.
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=DEVICE_LOCAL, usage=TRANSFER|DISPATCH_STORAGE, size=16
EXEC @multi_device_mul
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=DEVICE_LOCAL, usage=TRANSFER|DISPATCH_STORAGE, size=16
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=DEVICE_LOCAL, usage=TRANSFER|DISPATCH_STORAGE, size=64
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=DEVICE_LOCAL, usage=TRANSFER|DISPATCH_STORAGE, size=16
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=DEVICE_LOCAL, usage=TRANSFER|DISPATCH_STORAGE, size=64
iree_hal_metal_allocator_allocate_buffer
allocating metal buffer with params: type=HOST_LOCAL|DEVICE_VISIBLE, usage=TRANSFER|MAPPING, size=16
result[0]: hal.buffer_view
4xf32=0 55 144 273
https://gist.github.com/benvanik/ecc9b37fb2b670ce1ed2fb0d7c694287 shows what I'm shooting for after #20965 on the compiler side with zero copies once we have the transfer elision pass. Initially it'll let us do homogeneous zero-copy sharding as it sidesteps the runtime issues around synchronization. I'll see if I can get something quickly done for at least making the local executors able to wait and signal non-local semaphores - other targets will need to be able to wait on host events, though, via import or a waiter parking lot thread.
Dropping a note that https://github.com/iree-org/iree/blob/f26a830f71a35d3264333fa3f6a12f6cb5d35e30/compiler/src/iree/compiler/Dialect/Stream/Transforms/EmplaceAllocations.cpp#L91-L95 will need the topology check logic in order to perform in-place updates across devices.
@ziereis Im trying to build your work around and get to the point that you were. Trying to see if I can do any contribution here :)
Could you maybe include the toolchain that you used and your compile flags?
