fix transpose optimizer on GPU EP

[fs-eire]

Description

because of #15618 , the default allocator changed to device allocator, which will be GPU instead of CPU. in transpose optimizer we expect to read data from initializers so a CPU allocator is required here.

this change fixes transpose optimizer on GPU EP

Fixes the issue referred to in #15869, #15796

[askhade]

Looks like this PR need to be cherry picked as well.

[fs-eire]

Looks like this change fails tests on CUDA EP. I don't know much about how CUDA works with this transpose optimizer. Maybe I should only apply this fix to JSEP?

OK, I fixed the errors. So an EP may not register an allocator with type OrtMemTypeCPU (eg. InternalTestingExecutionProvider). So the latest behavior is, try to get allocator for OrtMemTypeDefault with that EP, if it is not on CPU device, try get allocator for OrtMemTypeCPU. If it still failed, return a failure status code.

[skottmckay]

Would it be simpler for the caller to pass in the value returned by ExecutionProviders.GetDefaultCpuAllocator()?

[fs-eire]

it's used in this way -

Status PartitionOrtFormatModel(onnxruntime::Graph& graph,
                               const ExecutionProviders& providers,
                               KernelRegistryManager& kernel_registry_manager,
                               SessionState& session_state) {

  layout_transformer::TransformLayoutFunction transform_layout_fn = layout_transformer::IsSupportedOpset(graph)
                                                                        ? layout_transformer::TransformLayoutForEP
                                                                        : nullptr;

  GraphPartitioner partitioner(kernel_registry_manager, providers);
  ORT_RETURN_IF_ERROR(partitioner.Partition(graph,
                                            session_state.GetMutableFuncMgr(),
                                            transform_layout_fn,
                                            GraphPartitioner::Mode::kOrtFormatLoad));

  return Status::OK();
}

So this is going to modify the function signature of GraphPartitioner::Partition() and some callback type definition, which is not a trivial one.

[skottmckay]

In InferenceSessions.cs there is a lambda used so technically you'd only need to update TransformLayoutForEP and not the entire path through GraphPartitioner.

onnxruntime/onnxruntime/core/session/inference_session.cc

Line 944 in 648bedf

layout_transformer::TransformLayoutForEP(graph_to_transform, modified, execution_provider,

[fs-eire]

There are 2 places where TransformLayoutForEP() is called, and neither of them has a context with an existing allocator. If I make this change, I still need code to get allocator from the EP, and I need do this twice.

Edit: I see you mentioned ExecutionProviders.GetDefaultCpuAllocator(). Let me check if I can figure it out how to do this

[fs-eire]

Updated the code

[fs-eire]

Is there an easier way to create lambdas in session_state_test.cc? is there a way to do currying in a simpler way