Add EP-specific weight layout transformation framework

This infrastructure enables execution providers to optimize operator
weights with custom memory layouts (such as blocked formats) during
session initialization, dramatically improving inference performance
through better cache utilization and memory access patterns.

Current Implementation:
- HWIO Transpose (WebGPU EP): Transposes Conv weights from OIHW to HWIO
layout as the first application of the framework
- ABcd16a4b Blocking (Proof-of-Concept): OneDNN-style blocked format
with 16×4 tiles demonstrates the framework's primary purpose

The framework is generic and extensible, allowing any EP to implement
custom weight transformations optimized for their target hardware.

Author: jchen10

include/onnxruntime/core/framework/execution_provider.h

@@ -378,6 +378,39 @@ class IExecutionProvider {
     return std::nullopt;
   }
 
+  /**
+    Query the preferred format descriptor for an initializer without performing the transformation.
+    This is a lightweight query called during session initialization to determine what format
+    transformations are needed.
+
+    @param node The node that consumes the initializer
+    @param input_index The input index of the initializer in the node
+    @param[out] format_descriptor A string that uniquely identifies the preferred format.
+                                  Empty string means no transformation is needed.
+                                  Examples: "ABcd16a4b", "hwio".
+    @return Status::OK() if query succeeded (format_descriptor will be set).
+            Failed status indicates no transformation is needed.
+  */
+  virtual Status GetPreferredInitializerFormat(const Node& /*node*/, int /*input_index*/,
+                                               std::string& /*format_descriptor*/) const {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "No format transformation needed");
+  }
+
+  /**
+    Transform an initializer to the specified format.
+    This performs the actual data transformation. It is only called once per unique format
+    even if multiple nodes need the same format.
+
+    @param original_tensor The original initializer tensor
+    @param format_descriptor The target format (from GetPreferredInitializerFormat)
+    @param[out] transformed_tensor The EP should allocate and fill this with the transformed data.
+    @return Status::OK() if transformation succeeded.
+  */
+  virtual Status TransformInitializerFormat(const Tensor& /*original_tensor*/, const std::string& /*format_descriptor*/,
+                                            std::unique_ptr<Tensor>& /*transformed_tensor*/) const {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Format transformation not supported");
+  }
+
   virtual void RegisterStreamHandlers(IStreamCommandHandleRegistry& /*stream_handle_registry*/, AllocatorMap&) const {}
 
   /** Does the EP support concurrent calls to InferenceSession::Run to execute the model.

include/onnxruntime/core/framework/tensor.h

@@ -282,6 +282,19 @@ class Tensor final {
     byte_offset_ = byte_offset;
   }
 
+  /**
+   * Get the memory format descriptor for this tensor.
+   * Returns empty string if the tensor is in standard format.
+   */
+  inline const std::string& GetFormatDescriptor() const { return format_descriptor_; }
+
+  /**
+   * Set the memory format descriptor for this tensor.
+   * Used for EP-specific memory layouts (e.g., "ABcd16a4b" for blocked format).
+   * The format string encodes all necessary information including block sizes.
+   */
+  inline void SetFormatDescriptor(const std::string& format) { format_descriptor_ = format; }
+
   /// <summary>
   /// The number of Tensor "storage" elements. A single storage element may contain multiple sub-elements for
   /// sub-byte data types (e.g., int4/float4).
@@ -349,6 +362,9 @@ class Tensor final {
   const PrimitiveDataTypeBase* dtype_;
   OrtMemoryInfo alloc_info_;
   ptrdiff_t byte_offset_;
+
+  // Memory format descriptor for EP-specific layouts (e.g., "ABcd16a4b")
+  std::string format_descriptor_;
 };
 #ifdef __GNUC__
 #pragma GCC diagnostic pop

onnxruntime/core/framework/session_state.cc

@@ -15,6 +15,7 @@
 #include "core/framework/ort_value_pattern_planner.h"
 #include "core/framework/prepacked_weights_container.h"
 #include "core/framework/session_state_utils.h"
+#include "core/framework/tensorprotoutils.h"
 #include "core/framework/utils.h"
 #include "core/providers/cpu/controlflow/utils.h"
 #include "core/session/onnxruntime_session_options_config_keys.h"
@@ -1332,6 +1333,9 @@ Status SessionState::FinalizeSessionState(const std::basic_string<PATH_CHAR_TYPE
   ORT_RETURN_IF_ERROR(CreateSubgraphSessionState());
 
   ORT_RETURN_IF_ERROR(VerifyEachNodeIsAssignedToAnEp(graph_, logger_, execution_providers_));
+
+  ORT_RETURN_IF_ERROR(TransformInitializersToPreferredFormat());
+
   ORT_RETURN_IF_ERROR(PopulateKernelCreateInfo(kernel_registry_manager, saving_ort_format));
 
   InlinedHashMap<std::string, size_t> constant_initializers_use_count;
@@ -1501,6 +1505,14 @@ Status SessionState::FinalizeSessionStateImpl(const std::basic_string<PATH_CHAR_
     CreateGraphInfo(save_prepacked_initializers);
   }
 
+  // Index all initializers including those that may have become unreferenced after transformation.
+  // This runs after CreateGraphInfo() to ensure consistent ordering - CreateGraphInfo indexes based on
+  // graph structure, then we add any remaining initializers (e.g., original weights before transformation).
+  for (const auto& [init_name, tensor_proto] : graph_.GetAllInitializedTensors()) {
+    ORT_UNUSED_PARAMETER(tensor_proto);
+    ort_value_name_idx_map_.Add(init_name);
+  }
+
 #if defined(ORT_EXTENDED_MINIMAL_BUILD)
   // Remove any unused initializers.
   // Not needed in a full build because unused initializers should have been removed earlier by Graph::Resolve().
@@ -1793,4 +1805,152 @@ void SessionState::RecycleDeviceStreamCollection(std::unique_ptr<DeviceStreamCol
 }
 #endif
 
+Status SessionState::TransformInitializersToPreferredFormat() {
+  // Build a map from initializer name to all nodes that consume it
+  std::unordered_map<std::string, std::vector<std::pair<NodeIndex, int>>> initializer_to_consumers;
+
+  const auto& initialized_tensors_map = graph_.GetAllInitializedTensors();
+  std::unordered_set<std::string> initializer_names;
+  for (const auto& [name, tensor_proto] : initialized_tensors_map) {
+    ORT_UNUSED_PARAMETER(tensor_proto);
+    initializer_names.insert(name);
+  }
+
+  // Scan nodes to find which initializers they use
+  for (const auto& node : graph_.Nodes()) {
+    int input_index = 0;
+    for (const auto* input_def : node.InputDefs()) {
+      if (input_def && input_def->Exists()) {
+        const auto& input_name = input_def->Name();
+        if (initializer_names.count(input_name) > 0) {
+          initializer_to_consumers[input_name].emplace_back(node.Index(), input_index);
+        }
+      }
+      ++input_index;
+    }
+  }
+
+  auto cpu_allocator = GetAllocator(OrtDevice());
+
+  for (const auto& [init_name, consumers] : initializer_to_consumers) {
+    if (consumers.empty()) {
+      continue;
+    }
+
+    const ONNX_NAMESPACE::TensorProto* tensor_proto = graph_.GetInitializer(init_name, true);
+    if (!tensor_proto) {
+      continue;
+    }
+
+    // Skip if this initializer was already transformed (when loading a saved ORT format model)
+    // Transformed initializers have format metadata in string_data
+    bool already_transformed = false;
+    for (const auto& attr_str : tensor_proto->string_data()) {
+      if (attr_str.find("onnxruntime_format:") == 0) {
+        already_transformed = true;
+        break;
+      }
+    }
+    if (already_transformed) {
+      continue;
+    }
+
+    // Phase 1: Query all consumers to discover what formats are needed
+    // Multiple nodes may request the same format, so we deduplicate by format
+    std::unordered_map<std::string, std::vector<std::pair<NodeIndex, int>>> format_to_consumers;
+
+    for (const auto& [node_idx, input_idx] : consumers) {
+      const Node* node = graph_.GetNode(node_idx);
+      if (!node) {
+        continue;
+      }
+
+      const auto& ep_type = node->GetExecutionProviderType();
+      if (ep_type.empty()) {
+        continue;
+      }
+
+      const auto* ep = execution_providers_.Get(ep_type);
+      if (!ep) {
+        continue;
+      }
+
+      // Ask EP if it wants this initializer in a different format
+      std::string format_descriptor;
+      Status query_status = ep->GetPreferredInitializerFormat(*node, input_idx, format_descriptor);
+
+      if (!query_status.IsOK() || format_descriptor.empty()) {
+        continue;
+      }
+
+      format_to_consumers[format_descriptor].emplace_back(node_idx, input_idx);
+    }
+
+    if (format_to_consumers.empty()) {
+      continue;
+    }
+
+    // Load the original initializer to CPU for transformation
+    TensorShape tensor_shape = utils::GetTensorShapeFromTensorProto(*tensor_proto);
+    const auto* tensor_type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto->data_type())->GetElementType();
+
+    Tensor original_tensor(tensor_type, tensor_shape, cpu_allocator);
+    ORT_RETURN_IF_ERROR(
+        utils::TensorProtoToTensor(Env::Default(), std::filesystem::path(), *tensor_proto, original_tensor));
+
+    // Phase 2: Transform once per unique format requested
+    for (const auto& [format_descriptor, nodes_needing_format] : format_to_consumers) {
+      const Node* first_node = graph_.GetNode(nodes_needing_format[0].first);
+      if (!first_node) {
+        continue;
+      }
+
+      const auto& ep_type = first_node->GetExecutionProviderType();
+      const auto* ep = execution_providers_.Get(ep_type);
+      if (!ep) {
+        continue;
+      }
+
+      // Perform the actual transformation
+      std::unique_ptr<Tensor> transformed_tensor;
+      Status transform_status = ep->TransformInitializerFormat(original_tensor, format_descriptor, transformed_tensor);
+
+      if (!transform_status.IsOK() || !transformed_tensor) {
+        LOGS(logger_, WARNING) << "Failed to transform initializer '" << init_name << "' to format '"
+                               << format_descriptor << "': " << transform_status.ErrorMessage();
+        continue;
+      }
+
+      // Set format metadata on the transformed tensor
+      transformed_tensor->SetFormatDescriptor(format_descriptor);
+
+      // Add the transformed initializer with a new name
+      std::string transformed_name = init_name + "_fmt_" + format_descriptor;
+
+      ONNX_NAMESPACE::TensorProto transformed_proto = utils::TensorToTensorProto(*transformed_tensor, transformed_name);
+
+      // Add format metadata as TensorProto attribute
+      auto* format_attr = transformed_proto.add_string_data();
+      *format_attr = "onnxruntime_format:" + format_descriptor;
+
+      graph_.AddInitializedTensor(transformed_proto);
+
+      // Update all nodes that need this format to use the transformed version
+      for (const auto& [node_idx, input_idx] : nodes_needing_format) {
+        Node* node = graph_.GetNode(node_idx);
+        if (!node) {
+          continue;
+        }
+
+        const auto* original_node_arg = node->InputDefs()[input_idx];
+        auto* transformed_node_arg = &graph_.GetOrCreateNodeArg(transformed_name, original_node_arg->TypeAsProto());
+
+        node->MutableInputDefs()[input_idx] = transformed_node_arg;
+      }
+    }
+  }
+
+  return Status::OK();
+}
+
 }  // namespace onnxruntime

onnxruntime/core/framework/session_state.h

@@ -431,6 +431,12 @@ class SessionState {
                                   const InlinedHashMap<OrtValueName, OrtDevice>& outer_scope_node_arg_to_location_map = {},
                                   bool graph_info_already_created = false);
 
+  /**
+   * Transform initializer tensors to EP-preferred memory formats.
+   * This is called during session initialization before kernel creation.
+   */
+  Status TransformInitializersToPreferredFormat();
+
 #ifdef ENABLE_TRAINING
   Status GeneratePatternGroupCache(
       gsl::span<const OrtValue> inputs,

onnxruntime/core/framework/session_state_utils.cc

@@ -227,6 +227,11 @@ common::Status CopyTensorFromCPUToDevice(
     }
     return copy_status;
   } else {
+    // Preserve format descriptor when copying from CPU to device
+    const std::string& format = deserialized_tensor.GetFormatDescriptor();
+    if (!format.empty()) {
+      tensor.SetFormatDescriptor(format);
+    }
     Tensor::InitOrtValue(std::move(tensor), ort_value);
     return common::Status::OK();
   }

onnxruntime/core/framework/tensor.cc

@@ -195,7 +195,8 @@ Tensor::Tensor(Tensor&& other) noexcept
 #endif
       dtype_(other.dtype_),
       alloc_info_(other.alloc_info_),
-      byte_offset_(other.byte_offset_) {
+      byte_offset_(other.byte_offset_),
+      format_descriptor_(std::move(other.format_descriptor_)) {
   other.p_data_ = nullptr;
   other.buffer_deleter_ = nullptr;
   other.dtype_ = DataTypeImpl::GetType<float>()->AsPrimitiveDataType();
@@ -221,6 +222,7 @@ Tensor& Tensor::operator=(Tensor&& other) noexcept {
     dtype_ = other.dtype_;
     alloc_info_ = other.alloc_info_;
     byte_offset_ = other.byte_offset_;
+    format_descriptor_ = std::move(other.format_descriptor_);
 
     other.p_data_ = nullptr;
     other.buffer_deleter_ = nullptr;

onnxruntime/core/framework/tensorprotoutils.cc

@@ -1424,6 +1424,15 @@ Status TensorProtoToTensor(const Env& env, const std::filesystem::path& model_pa
     }
   }
 
+  // Read format metadata from TensorProto string_data
+  for (const auto& attr_str : tensor_proto.string_data()) {
+    if (attr_str.find("onnxruntime_format:") == 0) {
+      std::string format = attr_str.substr(19);  // Skip "onnxruntime_format:"
+      tensor.SetFormatDescriptor(format);
+      break;  // Only one format descriptor expected
+    }
+  }
+
   return Status::OK();
 }