Merge remote-tracking branch 'origin/master' into edgchen1/merge_from_master

.gitignore

@@ -45,4 +45,4 @@ java/gradlew
 java/gradlew.bat
 java/gradle
 java/.gradle
-
+java/hs_*.log

docs/NotesOnThreading.md

@@ -0,0 +1,18 @@
+# Notes on Threading in ORT
+
+This document is intended for ORT developers.
+
+ORT allows the usage of either OpenMP or non-OpenMP (ORT) threads for execution. Threadpool management
+is abstracted behind: (1) ThreadPool class in threadpool.h and (2) functions in thread_utils.h.
+
+When developing an op, please use these abstractions to parallelize your code. These abstractions centralize 2 things.
+When OpenMP is enabled, they resort to using OpenMP. When OpenMP is disabled they resort to sequential execution if the threadpool ptr is NULL or schedule the tasks on the threadpool otherwise.
+
+Examples of these abstractions are: (threadpool.h has more documentation for these)
+* TryBatchParallelFor
+* TryParallelFor
+* static version of NumThreads
+
+**Please do not write #ifdef pragma omp in operator code**.
+
+For intra op parallelism ORT users can use either OpenMP or ORT threadpool. The choice of using OpenMP is indicated by building ORT with ```--use_openmp``` switch. For inter op parallelism, however, we always use the ORT threadpool.

include/onnxruntime/core/platform/threadpool.h

@@ -158,14 +158,12 @@ class ThreadPool {
                    const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& fn);
   static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, double cost_per_unit,
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& fn) {
-    if (tp == nullptr) {
-      fn(0, total);
-      return;
-    }
-    tp->ParallelFor(total, cost_per_unit, fn);
+    TryParallelFor(tp, total, TensorOpCost{0, 0, static_cast<double>(cost_per_unit)}, fn);
   }
+
   void ParallelFor(std::ptrdiff_t total, const TensorOpCost& cost_per_unit,
                    const std::function<void(std::ptrdiff_t first, std::ptrdiff_t)>& fn);
+
   static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const TensorOpCost& cost_per_unit,
                              const std::function<void(std::ptrdiff_t first, std::ptrdiff_t last)>& fn) {
     if (tp == nullptr) {
@@ -174,10 +172,12 @@ class ThreadPool {
     }
     tp->ParallelFor(total, cost_per_unit, fn);
   }
+
   // Similar to ParallelFor above, but takes the specified scheduling strategy
   // into account.
-  void ParallelFor(std::ptrdiff_t total, const SchedulingParams& scheduling_params,
-                   const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn);
+  void
+  ParallelFor(std::ptrdiff_t total, const SchedulingParams& scheduling_params,
+              const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn);
 
   static void TryParallelFor(concurrency::ThreadPool* tp, std::ptrdiff_t total, const SchedulingParams& scheduling_params,
                              const std::function<void(std::ptrdiff_t, std::ptrdiff_t)>& fn) {
@@ -187,7 +187,12 @@ class ThreadPool {
     }
     tp->ParallelFor(total, scheduling_params, fn);
   }
-  // Returns the number of threads in the pool.
+
+  // Prefer using this API to get the number of threads unless you know what you're doing.
+  // This API takes into account if openmp is enabled/disabled and if the thread pool ptr is nullptr.
+  static int NumThreads(const concurrency::ThreadPool* tp);
+
+  // Returns the number of threads in the pool. Preferably use the static version of this API instead.
   int NumThreads() const;
 
   // Returns current thread id between 0 and NumThreads() - 1, if called from a

java/build.gradle

@@ -1,6 +1,7 @@
 plugins {
 	id 'java'
 	id 'jacoco'
+	id 'maven-publish'
 	id 'com.diffplug.gradle.spotless' version '3.26.0'
 }
 
@@ -13,8 +14,36 @@ allprojects {
 java {
 	sourceCompatibility = JavaVersion.VERSION_1_8
 	targetCompatibility = JavaVersion.VERSION_1_8
-	withJavadocJar()
-	withSourcesJar()
+}
+
+project.group = "ai.onnxruntime"
+version = rootProject.file('../VERSION_NUMBER').text.trim()
+
+jar {
+	into("META-INF/maven/$project.group/$project.name") {
+		from { generatePomFileForMavenPublication }
+		rename ".*", "pom.xml"
+	}
+}
+
+// Add explicit sources jar with pom file.
+task sourcesJar(type: Jar, dependsOn: classes) {
+	classifier = "sources"
+	from sourceSets.main.allSource
+	into("META-INF/maven/$project.group/$project.name") {
+		from { generatePomFileForMavenPublication }
+		rename ".*", "pom.xml"
+	}
+}
+
+// Add explicit javadoc jar with pom file
+task javadocJar(type: Jar, dependsOn: javadoc) {
+	classifier = "javadoc"
+	from javadoc.destinationDir
+	into("META-INF/maven/$project.group/$project.name") {
+		from { generatePomFileForMavenPublication }
+		rename ".*", "pom.xml"
+	}
 }
 
 wrapper {
@@ -40,8 +69,6 @@ def cmakeNativeLibDir = "${cmakeJavaDir}/native-lib"
 def cmakeNativeJniDir = "${cmakeJavaDir}/native-jni"
 def cmakeBuildOutputDir = "${cmakeJavaDir}/build"
 
-version = rootProject.file('../VERSION_NUMBER').text.trim()
-
 compileJava {
 	options.compilerArgs += ["-h", "${project.buildDir}/headers/"]
 }
@@ -76,6 +103,10 @@ if (cmakeBuildDir != null) {
 	}
 
 	task allJar(type: Jar) {
+		into("META-INF/maven/$project.group/$project.name") {
+			from { generatePomFileForMavenPublication }
+			rename ".*", "pom.xml"
+		}
 		classifier = 'all'
 		from sourceSets.main.output
 		from cmakeNativeJniDir
@@ -96,7 +127,6 @@ if (cmakeBuildDir != null) {
 	cmakeBuild.dependsOn javadocJar
 	cmakeBuild.dependsOn javadoc
 
-
 	task cmakeCheck(type: Copy) {
 		from project.buildDir
 		include 'reports/**'
@@ -116,6 +146,7 @@ test {
 	useJUnitPlatform()
 	testLogging {
 		events "passed", "skipped", "failed"
+		showStandardStreams = true
 	}
 }
 
@@ -126,3 +157,34 @@ jacocoTestReport {
 		html.destination file("${buildDir}/jacocoHtml")
 	}
 }
+
+publishing {
+	publications {
+		maven(MavenPublication) {
+			groupId = project.group
+			artifactId = project.name
+
+			from components.java
+			pom {
+				name = 'onnx-runtime'
+				description = 'ONNX Runtime is a performance-focused inference engine for ONNX (Open Neural Network Exchange) models.'
+				url = 'https://microsoft.github.io/onnxruntime/'
+				licenses {
+					license {
+						name = 'MIT License'
+						url = 'https://opensource.org/licenses/MIT'
+					}
+				}
+				organization {
+					name = 'Microsoft'
+					url = 'http://www.microsoft.com'
+				}
+				scm {
+					connection = 'scm:git:git://github.com:microsoft/onnxruntime.git'
+					developerConnection = 'scm:git:ssh://github.com/microsoft/onnxruntime.git'
+					url = 'http://github.com/microsoft/onnxruntime'
+				}
+			}
+		}
+	}
+}

java/src/main/java/ai/onnxruntime/OnnxModelMetadata.java

@@ -0,0 +1,210 @@
+package ai.onnxruntime;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+
+/**
+ * Contains the metadata associated with an ONNX model.
+ *
+ * <p>Unspecified default fields contain the empty string.
+ *
+ * <p>This class is a Java side copy of the native metadata, it does not access the native runtime.
+ */
+public final class OnnxModelMetadata {
+
+  private final String producerName;
+  private final String graphName;
+  private final String domain;
+  private final String description;
+  private final long version;
+
+  private final Map<String, String> customMetadata;
+
+  /**
+   * Constructed by an OrtSession in native code, nulls are replaced with the empty string or empty
+   * map as appropriate.
+   *
+   * @param producerName The model producer name.
+   * @param graphName The model graph name.
+   * @param domain The model domain name.
+   * @param description The model description.
+   * @param version The model version.
+   * @param customMetadataArray Any custom metadata associated with the model.
+   */
+  OnnxModelMetadata(
+      String producerName,
+      String graphName,
+      String domain,
+      String description,
+      long version,
+      String[] customMetadataArray) {
+    this.producerName = producerName == null ? "" : producerName;
+    this.graphName = graphName == null ? "" : graphName;
+    this.domain = domain == null ? "" : domain;
+    this.description = description == null ? "" : description;
+    this.version = version;
+    if (customMetadataArray != null && customMetadataArray.length > 0) {
+      this.customMetadata = new HashMap<>();
+      if (customMetadataArray.length % 2 == 1) {
+        throw new IllegalStateException(
+            "Asked for keys and values, but received an odd number of elements.");
+      }
+      for (int i = 0; i < customMetadataArray.length; i += 2) {
+        customMetadata.put(customMetadataArray[i], customMetadataArray[i + 1]);
+      }
+    } else {
+      this.customMetadata = Collections.emptyMap();
+    }
+  }
+
+  /**
+   * Constructed by an OrtSession, nulls are replaced with the empty string or empty map as
+   * appropriate.
+   *
+   * @param producerName The model producer name.
+   * @param graphName The model graph name.
+   * @param domain The model domain name.
+   * @param description The model description.
+   * @param version The model version.
+   * @param customMetadata Any custom metadata associated with the model.
+   */
+  OnnxModelMetadata(
+      String producerName,
+      String graphName,
+      String domain,
+      String description,
+      long version,
+      Map<String, String> customMetadata) {
+    this.producerName = producerName == null ? "" : producerName;
+    this.graphName = graphName == null ? "" : graphName;
+    this.domain = domain == null ? "" : domain;
+    this.description = description == null ? "" : description;
+    this.version = version;
+    this.customMetadata = customMetadata == null ? Collections.emptyMap() : customMetadata;
+  }
+
+  /**
+   * Copy constructor.
+   *
+   * @param other The metadata to copy.
+   */
+  public OnnxModelMetadata(OnnxModelMetadata other) {
+    this.producerName = other.producerName;
+    this.graphName = other.graphName;
+    this.domain = other.domain;
+    this.description = other.description;
+    this.version = other.version;
+    this.customMetadata =
+        other.customMetadata.isEmpty()
+            ? Collections.emptyMap()
+            : new HashMap<>(getCustomMetadata());
+  }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    OnnxModelMetadata that = (OnnxModelMetadata) o;
+    return version == that.version
+        && producerName.equals(that.producerName)
+        && graphName.equals(that.graphName)
+        && domain.equals(that.domain)
+        && description.equals(that.description)
+        && customMetadata.equals(that.customMetadata);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(producerName, graphName, domain, description, version, customMetadata);
+  }
+
+  /**
+   * Gets the producer name.
+   *
+   * @return The producer name.
+   */
+  public String getProducerName() {
+    return producerName;
+  }
+
+  /**
+   * Gets the graph name.
+   *
+   * @return The graph name.
+   */
+  public String getGraphName() {
+    return graphName;
+  }
+
+  /**
+   * Gets the domain.
+   *
+   * @return The domain.
+   */
+  public String getDomain() {
+    return domain;
+  }
+
+  /**
+   * Gets the model description.
+   *
+   * @return The description.
+   */
+  public String getDescription() {
+    return description;
+  }
+
+  /**
+   * Gets the model version.
+   *
+   * @return The model version.
+   */
+  public long getVersion() {
+    return version;
+  }
+
+  /**
+   * Gets an unmodifiable reference to the complete custom metadata.
+   *
+   * @return The custom metadata.
+   */
+  public Map<String, String> getCustomMetadata() {
+    return Collections.unmodifiableMap(customMetadata);
+  }
+
+  /**
+   * Returns Optional.of(value) if the custom metadata has a value for the supplied key, otherwise
+   * returns {@link Optional#empty}.
+   *
+   * @param key The custom metadata key.
+   * @return The custom metadata value if present.
+   */
+  public Optional<String> getCustomMetadataValue(String key) {
+    return Optional.ofNullable(customMetadata.get(key));
+  }
+
+  @Override
+  public String toString() {
+    return "OnnxModelMetadata{"
+        + "producerName='"
+        + producerName
+        + '\''
+        + ", graphName='"
+        + graphName
+        + '\''
+        + ", domain='"
+        + domain
+        + '\''
+        + ", description='"
+        + description
+        + '\''
+        + ", version="
+        + version
+        + ", customMetadata="
+        + customMetadata
+        + '}';
+  }
+}

java/src/main/java/ai/onnxruntime/OnnxRuntime.java

@@ -24,6 +24,8 @@ final class OnnxRuntime {
   private static final int ORT_API_VERSION_1 = 1;
   // Post 1.0 builds of the ORT API.
   private static final int ORT_API_VERSION_2 = 2;
+  // Post 1.3 builds of the ORT API
+  private static final int ORT_API_VERSION_3 = 3;
 
   /** The short name of the ONNX runtime shared library */
   static final String ONNXRUNTIME_LIBRARY_NAME = "onnxruntime";
@@ -50,7 +52,7 @@ final class OnnxRuntime {
     try {
       load(tempDirectory, ONNXRUNTIME_LIBRARY_NAME);
       load(tempDirectory, ONNXRUNTIME_JNI_LIBRARY_NAME);
-      ortApiHandle = initialiseAPIBase(ORT_API_VERSION_2);
+      ortApiHandle = initialiseAPIBase(ORT_API_VERSION_3);
       loaded = true;
     } finally {
       if (!isAndroid()) {

java/src/main/java/ai/onnxruntime/OrtSession.java

@@ -47,6 +47,8 @@ public class OrtSession implements AutoCloseable {
 
   private final Set<String> outputNames;
 
+  private OnnxModelMetadata metadata;
+
   private boolean closed = false;
 
   /**
@@ -201,7 +203,7 @@ public class OrtSession implements AutoCloseable {
   /**
    * Scores an input feed dict, returning the map of requested inferred outputs.
    *
-   * <p>The outputs are sorted based on the supplied set traveral order.
+   * <p>The outputs are sorted based on the supplied set traversal order.
    *
    * @param inputs The inputs to score.
    * @param requestedOutputs The requested outputs.
@@ -263,6 +265,18 @@ public class OrtSession implements AutoCloseable {
     }
   }
 
+  /**
+   * Gets the metadata for the currently loaded model.
+   *
+   * @return The metadata.
+   */
+  public OnnxModelMetadata getMetadata() throws OrtException {
+    if (metadata == null) {
+      metadata = constructMetadata(OnnxRuntime.ortApiHandle, nativeHandle, allocator.handle);
+    }
+    return metadata;
+  }
+
   @Override
   public String toString() {
     return "OrtSession(numInputs=" + numInputs + ",numOutputs=" + numOutputs + ")";
@@ -334,6 +348,18 @@ public class OrtSession implements AutoCloseable {
 
   private native void closeSession(long apiHandle, long nativeHandle) throws OrtException;
 
+  /**
+   * Builds the {@link OnnxModelMetadata} for this session.
+   *
+   * @param ortApiHandle The api pointer.
+   * @param nativeHandle The native session pointer.
+   * @param allocatorHandle The OrtAllocator pointer.
+   * @return The metadata.
+   * @throws OrtException If the native runtime failed to access or allocate the metadata.
+   */
+  private native OnnxModelMetadata constructMetadata(
+      long ortApiHandle, long nativeHandle, long allocatorHandle) throws OrtException;
+
   /**
    * Represents the options used to construct this session.
    *

java/src/main/native/ai_onnxruntime_OrtSession.c

@@ -337,3 +337,94 @@ JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_releaseNamesHandle
     }
     checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,names));
 }
+
+/*
+ * Class:     ai_onnxruntime_OrtSession
+ * Method:    constructMetadata
+ * Signature: (JJJ)Ljava/lang/String;
+ */
+JNIEXPORT jstring JNICALL Java_ai_onnxruntime_OrtSession_constructMetadata
+    (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong nativeHandle, jlong allocatorHandle) {
+  (void) jobj; // Required JNI parameter not needed by functions which don't need to access their host object.
+  const OrtApi* api = (const OrtApi*) apiHandle;
+  OrtAllocator* allocator = (OrtAllocator*) allocatorHandle;
+
+  // Setup
+  char* stringClassName = "java/lang/String";
+  jclass stringClazz = (*jniEnv)->FindClass(jniEnv, stringClassName);
+  char *metadataClassName = "ai/onnxruntime/OnnxModelMetadata";
+  jclass metadataClazz = (*jniEnv)->FindClass(jniEnv, metadataClassName);
+  //OnnxModelMetadata(String producerName, String graphName, String domain, String description, long version, String[] customMetadataArray)
+  jmethodID metadataConstructor = (*jniEnv)->GetMethodID(jniEnv, metadataClazz, "<init>",
+                                                         "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;J[Ljava/lang/String;)V");
+
+  // Get metadata
+  OrtModelMetadata* metadata;
+  checkOrtStatus(jniEnv,api,api->SessionGetModelMetadata((OrtSession*)nativeHandle,&metadata));
+
+  // Read out the producer name and convert it to a java.lang.String
+  char* charBuffer;
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetProducerName(metadata, allocator, &charBuffer));
+  jstring producerStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
+  checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
+
+  // Read out the graph name and convert it to a java.lang.String
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetGraphName(metadata, allocator, &charBuffer));
+  jstring graphStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
+  checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
+
+  // Read out the domain and convert it to a java.lang.String
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetDomain(metadata, allocator, &charBuffer));
+  jstring domainStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
+  checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
+
+  // Read out the description and convert it to a java.lang.String
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetDescription(metadata, allocator, &charBuffer));
+  jstring descriptionStr = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
+  checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
+
+  // Read out the version
+  int64_t version;
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetVersion(metadata, &version));
+
+  // Read out the keys, look up the values.
+  int64_t numKeys;
+  char** keys;
+  checkOrtStatus(jniEnv,api,api->ModelMetadataGetCustomMetadataMapKeys(metadata, allocator, &keys, &numKeys));
+  jobjectArray customArray = NULL;
+  if (numKeys > 0) {
+    customArray = (*jniEnv)->NewObjectArray(jniEnv,numKeys*2,stringClazz,NULL);
+
+    // Iterate key array to extract the values
+    for (int64_t i = 0; i < numKeys; i++) {
+      // Create a java.lang.String for the key
+      jstring keyJava = (*jniEnv)->NewStringUTF(jniEnv,keys[i]);
+
+      // Extract the value and convert it to a java.lang.String
+      checkOrtStatus(jniEnv,api,api->ModelMetadataLookupCustomMetadataMap(metadata,allocator,keys[i],&charBuffer));
+      jstring valueJava = (*jniEnv)->NewStringUTF(jniEnv,charBuffer);
+      checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,charBuffer));
+
+      // Write the key and value into the array
+      (*jniEnv)->SetObjectArrayElement(jniEnv, customArray, i*2, keyJava);
+      (*jniEnv)->SetObjectArrayElement(jniEnv, customArray, (i*2)+1, valueJava);
+    }
+
+    // Release key array
+    for (int64_t i = 0; i < numKeys; i++) {
+      checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,keys[i]));
+    }
+    checkOrtStatus(jniEnv,api,api->AllocatorFree(allocator,keys));
+  } else {
+    customArray = (*jniEnv)->NewObjectArray(jniEnv,0,stringClazz,NULL);
+  }
+
+  // Invoke the metadata constructor
+  //OnnxModelMetadata(String producerName, String graphName, String domain, String description, long version, String[] customMetadataArray)
+  jobject metadataJava = (*jniEnv)->NewObject(jniEnv, metadataClazz, metadataConstructor, producerStr, graphStr, domainStr, descriptionStr, (jlong) version, customArray);
+
+  // Release the metadata
+  api->ReleaseModelMetadata(metadata);
+
+  return metadataJava;
+}

java/src/test/java/ai/onnxruntime/InferenceTest.java

@@ -41,6 +41,8 @@ import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BiFunction;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
@@ -63,6 +65,7 @@ public class InferenceTest {
 
   @Test
   public void repeatedCloseTest() throws OrtException {
+    Logger.getLogger(OrtEnvironment.class.getName()).setLevel(Level.SEVERE);
     OrtEnvironment env = OrtEnvironment.getEnvironment("repeatedCloseTest");
     try (OrtEnvironment otherEnv = OrtEnvironment.getEnvironment()) {
       assertFalse(otherEnv.isClosed());
@@ -115,7 +118,9 @@ public class InferenceTest {
   @Test
   public void morePartialInputsTest() throws OrtException {
     String modelPath = getResourcePath("/partial-inputs-test-2.onnx").toString();
-    try (OrtEnvironment env = OrtEnvironment.getEnvironment("partialInputs");
+    try (OrtEnvironment env =
+            OrtEnvironment.getEnvironment(
+                OrtEnvironment.LoggingLevel.ORT_LOGGING_LEVEL_FATAL, "partialInputs");
         OrtSession.SessionOptions options = new SessionOptions();
         OrtSession session = env.createSession(modelPath, options)) {
       assertNotNull(session);
@@ -200,7 +205,9 @@ public class InferenceTest {
   @Test
   public void partialInputsTest() throws OrtException {
     String modelPath = getResourcePath("/partial-inputs-test.onnx").toString();
-    try (OrtEnvironment env = OrtEnvironment.getEnvironment("partialInputs");
+    try (OrtEnvironment env =
+            OrtEnvironment.getEnvironment(
+                OrtEnvironment.LoggingLevel.ORT_LOGGING_LEVEL_FATAL, "partialInputs");
         OrtSession.SessionOptions options = new SessionOptions();
         OrtSession session = env.createSession(modelPath, options)) {
       assertNotNull(session);
@@ -370,6 +377,15 @@ public class InferenceTest {
         for (int i = 0; i < expectedOutputDimensions.length; i++) {
           assertEquals(expectedOutputDimensions[i], outputInfo.shape[i]);
         }
+
+        // Check the metadata can be extracted
+        OnnxModelMetadata metadata = session.getMetadata();
+        assertEquals("onnx-caffe2", metadata.getProducerName());
+        assertEquals("squeezenet_old", metadata.getGraphName());
+        assertEquals("", metadata.getDomain());
+        assertEquals("", metadata.getDescription());
+        assertEquals(0x7FFFFFFFFFFFFFFFL, metadata.getVersion());
+        assertTrue(metadata.getCustomMetadata().isEmpty());
       }
     }
   }

onnxruntime/core/common/threadpool.cc

@@ -298,6 +298,15 @@ void ThreadPool::ParallelFor(std::ptrdiff_t total, double cost_per_unit,
   ParallelFor(total, TensorOpCost{0, 0, static_cast<double>(cost_per_unit)}, fn);
 }
 
+int ThreadPool::NumThreads(const concurrency::ThreadPool* tp) {
+#ifdef _OPENMP
+  ORT_UNUSED_PARAMETER(tp);
+  return (omp_get_num_threads() == 1) ? omp_get_max_threads() : 1;
+#else
+  return tp ? tp->NumThreads() : 1;
+#endif
+}
+
 int ThreadPool::NumThreads() const {
   return underlying_threadpool_->NumThreads();
 }

onnxruntime/core/mlas/lib/mlasi.h

@@ -97,6 +97,10 @@ Abstract:
 //
 // Select the threading model.
 //
+// N.B. MLAS_NO_ONNXRUNTIME_THREADPOOL is used to build MLAS test code outside
+// of the ONNX Runtime source tree. OpenMP may or may not be enabled in this
+// configuration.
+//
 
 #if !defined(MLAS_NO_ONNXRUNTIME_THREADPOOL)
 #include "core/platform/threadpool.h"
@@ -666,19 +670,17 @@ MlasGetMaximumThreadCount(
     MLAS_THREADPOOL* ThreadPool
     )
 {
-#ifdef MLAS_NO_ONNXRUNTIME_THREADPOOL
+#if defined(MLAS_NO_ONNXRUNTIME_THREADPOOL)
     MLAS_UNREFERENCED_PARAMETER(ThreadPool);
-#else
-    if (ThreadPool != nullptr) {
-        return ThreadPool->NumThreads();
-    }
-#endif
 
 #if defined(_OPENMP)
     return (omp_get_num_threads() == 1) ? omp_get_max_threads() : 1;
 #else
     return 1;
 #endif
+#else
+	return onnxruntime::concurrency::ThreadPool::NumThreads(ThreadPool);
+#endif
 }
 
 inline

onnxruntime/core/platform/env.h

@@ -113,6 +113,9 @@ class Env {
 
   virtual int GetNumCpuCores() const = 0;
 
+  // This function doesn't support systems with more than 64 logical processors
+  virtual std::vector<size_t> GetThreadAffinityMasks() const = 0;
+
   /// \brief Returns the number of micro-seconds since the Unix epoch.
   virtual uint64_t NowMicros() const {
     return env_time_->NowMicros();

onnxruntime/core/platform/posix/env.cc

@@ -141,7 +141,7 @@ class PosixThread : public EnvThread {
       if (s != 0)
         ORT_THROW("pthread_setaffinity_np failed");
     }
- #endif
+#endif
   }
 
   ~PosixThread() override {
@@ -199,6 +199,12 @@ class PosixEnv : public Env {
     return std::thread::hardware_concurrency();
   }
 
+  std::vector<size_t> GetThreadAffinityMasks() const override {
+    std::vector<size_t> ret(std::thread::hardware_concurrency() / 2);
+    std::iota(ret.begin(), ret.end(), 0);
+    return ret;
+  }
+
   void SleepForMicroseconds(int64_t micros) const override {
     while (micros > 0) {
       timespec sleep_time;

onnxruntime/core/platform/windows/env.cc

@@ -143,6 +143,30 @@ class WindowsEnv : public Env {
     return processorCoreCount;
   }
 
+  std::vector<size_t> GetThreadAffinityMasks() const override {
+    auto generate_vector_of_n = [](int n) {
+      std::vector<size_t> ret(n);
+      std::iota(ret.begin(), ret.end(), 0);
+      return ret;
+    };
+    // Indeed 64 should be enough. However, it's harmless to have a little more.
+    SYSTEM_LOGICAL_PROCESSOR_INFORMATION buffer[256];
+    DWORD returnLength = sizeof(buffer);
+    if (GetLogicalProcessorInformation(buffer, &returnLength) == FALSE) {
+      return generate_vector_of_n(std::thread::hardware_concurrency());
+    }
+    std::vector<size_t> ret;
+    int count = (int)(returnLength / sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION));
+    for (int i = 0; i != count; ++i) {
+      if (buffer[i].Relationship == RelationProcessorCore) {
+        ret.push_back(buffer[i].ProcessorMask);
+      }
+    }
+    if (ret.empty())
+      return generate_vector_of_n(std::thread::hardware_concurrency());
+    return ret;
+  }
+
   static WindowsEnv& Instance() {
     static WindowsEnv default_env;
     return default_env;

onnxruntime/core/providers/cpu/math/top_k.cc

@@ -20,6 +20,7 @@
 #include "core/common/exceptions.h"
 #include "core/framework/op_kernel.h"
 #include "core/framework/tensor.h"
+#include "core/platform/threadpool.h"
 #include "core/util/math_cpuonly.h"
 #include <queue>
 #include <algorithm>
@@ -31,145 +32,301 @@ namespace onnxruntime {
 template <typename T>
 struct GreaterValueCmp {
   using DataType = T;
-  bool operator()(const pair<T, int64_t>& lhs, const pair<T, int64_t>& rhs) {
-    return (lhs.first > rhs.first ||
+  GreaterValueCmp(const T* data = nullptr) : data_(data) {
+  }
+
+  bool operator()(const int64_t lhs_idx, const int64_t rhs_idx) const {
+    return (data_[lhs_idx] > data_[rhs_idx] ||
             // when values are equal, we want lhs to get higher "priority"
             // if its corresponding index comes first (i.e.) is lower
-            (lhs.first == rhs.first && lhs.second < rhs.second));
+            (data_[lhs_idx] == data_[rhs_idx] && lhs_idx < rhs_idx));
   }
+
+  bool CompareValueOnly(const T& lhs, const T& rhs) const {
+    return lhs > rhs;
+  }
+
+ private:
+  const T* data_;
 };
 
 template <typename T>
 struct LesserValueCmp {
   using DataType = T;
-  bool operator()(const pair<T, int64_t>& lhs, const pair<T, int64_t>& rhs) {
-    return (lhs.first < rhs.first ||
+
+  LesserValueCmp(const T* data = nullptr) : data_(data) {
+  }
+
+  bool operator()(const int64_t lhs_idx, const int64_t rhs_idx) const {
+    return (data_[lhs_idx] < data_[rhs_idx] ||
             // when values are equal, we want lhs to get higher "priority"
             // if its corresponding index comes first (i.e.) is lower
-            (lhs.first == rhs.first && lhs.second < rhs.second));
+            (data_[lhs_idx] == data_[rhs_idx] && lhs_idx < rhs_idx));
   }
+
+  bool CompareValueOnly(const T& lhs, const T& rhs) const {
+    return lhs < rhs;
+  }
+
+ private:
+  const T* data_;
 };
 
+/*
+Maintain a binary heap where HeapComp of the parent with either child is false.
+  e.g. if the comparison is 'greater than', the parent is smaller than both children.
+There is no ordering within a level.
+
+NOTE: The comparison is backwards compared to std::priority_queue as we use the same comparator for this as for
+      nth_element in SelectTopK. As such for a heap selecting the largest values the comparator is 'greater than'.
+*/
+template <class HeapCmp>
+static void HeapifyIthPosition(int64_t* heap, size_t i, size_t k, const HeapCmp& heap_cmp) {
+  while (true) {
+    size_t left = 2 * i + 1;
+    size_t right = left + 1;
+    if (right < k) {
+      // need to check both left and right children as either could be replaced
+
+      // check if we should move child up. check left node as well as whether left is preferred over right.
+      // if 'i' can replace left, check whether right would replace left (if so, i replaces left as it's the weakest)
+      bool i_replaces_left = heap_cmp(heap[i], heap[left]);
+      if (i_replaces_left && heap_cmp(heap[right], heap[left])) {
+        // left is going to be pushed up as both i and right beat it
+        // NOTE: std::swap is slower as it uses std::move
+        auto tmp = heap[i];
+        heap[i] = heap[left];
+        heap[left] = tmp;
+        i = left;
+      } else if (i_replaces_left || heap_cmp(heap[i], heap[right])) {
+        // i_replaces_left implies left replaces right due to 'if' so replace right with i as right is the weakest.
+        // also check if i only beats right
+        auto tmp = heap[i];
+        heap[i] = heap[right];
+        heap[right] = tmp;
+        i = right;
+      } else
+        break;
+    } else if ((left < k) && heap_cmp(heap[i], heap[left])) {
+      auto tmp = heap[i];
+      heap[i] = heap[left];
+      heap[left] = tmp;
+      i = left;
+    } else
+      break;
+  }
+}
+
 // Static helpers that implement the core logic for each of the 'TopK' operator flavor
 
 // Selects the top k elements (largest or smallest based on template parameter)
 template <class Comparator>
-static vector<pair<typename Comparator::DataType, int64_t>> select_top_k(
-    const ConstEigenMatrixMapRowMajor<typename Comparator::DataType>& raw_data, int64_t row_num, int64_t num_blocks,
-    int64_t block_slice, int64_t inter_block_offset, const unsigned k,
-    bool sort_top_k) {
-  // create a data holder and insert elements
-  vector<pair<typename Comparator::DataType, int64_t>> data_holder;
-  data_holder.reserve(num_blocks);
+static void SelectTopK(const Comparator& comparer,
+                       int64_t row_offset, int64_t num_blocks, int64_t block_slice, int64_t inter_block_offset,
+                       const unsigned k, bool sort_top_k, vector<int64_t>& data_holder) {
   for (int64_t l = 0; l < num_blocks; ++l) {
-    data_holder.push_back({raw_data(row_num, l * block_slice + inter_block_offset), l});
+    data_holder[l] = (row_offset + (l * block_slice + inter_block_offset));
   }
 
-  // find the top k (largest or smallest) elements in the data holder - O(n)
-  nth_element(data_holder.begin(), data_holder.begin() + (k - 1), data_holder.end(), Comparator());
+  // find the top k (largest or smallest) elements in the data holder - O(n) average. O(n*n) worst case.
+  // See https://en.wikipedia.org/wiki/Quickselect
+  nth_element(data_holder.begin(), data_holder.begin() + (k - 1), data_holder.end(), comparer);
 
   // sort the top k elements if needed - O (k log k)
   if (sort_top_k) {
-    std::sort(data_holder.begin(), data_holder.begin() + k, Comparator());
+    std::sort(data_holder.begin(), data_holder.begin() + k, comparer);
   }
 
-  // the data_holder now contains the top k elements in the first k indices
-  return data_holder;
+  // the data_holder now contains the indices of the top k elements in the first k elements
 }
 
 // Given an input tensor 'input' and metadata values - 'k' and 'axis_parsed',
 // this method will extract the sorted top k largest/smallest elements and place them in the output tensor 'values'
 // along with the metadata output 'indices'
-template <bool largest, bool sorted, class Comparator>
-static void extract_top_k_elements(const Tensor* input, const TensorShape& input_shape, Tensor* values,
-                                   Tensor* indices, const TensorShape& output_shape, const unsigned k,
-                                   const unsigned axis_parsed) {
+template <class Comparator>
+static void FindTopKElements(const Tensor* input, const TensorShape& input_shape, Tensor* values,
+                             Tensor* indices, const TensorShape& output_shape, const unsigned k, bool sorted,
+                             const unsigned axis_parsed, concurrency::ThreadPool* threadpool) {
   // Cache some values that will be used in the implementation below
   const int64_t rows = input_shape.SizeToDimension(static_cast<size_t>(axis_parsed));
   const int64_t cols = input->Shape().Size() / rows;
-  auto input_map =
-      ConstEigenMatrixMapRowMajor<typename Comparator::DataType>(
-          static_cast<const typename Comparator::DataType*>(input->template Data<typename Comparator::DataType>()), rows, cols);
+  const auto* input_data = input->template Data<typename Comparator::DataType>();
 
-  // Use Eigen maps to allow indexing into the 2d tensors like Values_map(i,j)
+  // Use Eigen maps for convenient indexing into the 2d tensors like Values_map(i,j)
   const int64_t reduced_cols = output_shape.SizeFromDimension(static_cast<size_t>(axis_parsed));
-  auto values_map = EigenMatrixMapRowMajor<typename Comparator::DataType>(
-      values->template MutableData<typename Comparator::DataType>(), rows, reduced_cols);
-  auto indices_map = EigenMatrixMapRowMajor<int64_t>(indices->template MutableData<int64_t>(), rows, reduced_cols);
+
+  auto* values_data = values->template MutableData<typename Comparator::DataType>();
+  auto* indices_data = indices->template MutableData<int64_t>();
+  auto values_map = EigenMatrixMapRowMajor<typename Comparator::DataType>(values_data, rows, reduced_cols);
+  auto indices_map = EigenMatrixMapRowMajor<int64_t>(indices_data, rows, reduced_cols);
 
   // This is basically the number of elements within each of the "k" rows
-  const int64_t block_slice = reduced_cols / k;
   const int64_t num_blocks = input_shape[axis_parsed];
+  const int64_t block_slice = reduced_cols / k;
 
-  for (int64_t i = 0; i < rows; ++i) {
-    for (int64_t j = 0; j < block_slice; ++j) {
-      // Since sorted == true, we will use a Heap to hold the top K values in sorted fashion
-      if (sorted) {  // The optimizer will clean-up the redundant condition based on the template parameter 'sorted'
-        auto n_casted = static_cast<double>(num_blocks);
-        auto k_casted = static_cast<double>(k);
-        if ((n_casted + k_casted * log(k_casted)) < (n_casted * log(k_casted))) {
-          // Select first  - O(n), then sort O(k * ln(k))
-          // Overall complexity =  O (n + k * ln(k))
-          const auto& data_holder = select_top_k<Comparator>(input_map, i, num_blocks, block_slice, j, k, true);
-          for (int64_t l = 0; l < k; ++l) {
-            const auto& elem = data_holder[l];
-            auto col_index = l * block_slice + j;
-            values_map(i, col_index) = elem.first;
-            indices_map(i, col_index) = elem.second;
-          }
-        } else {
-          // Perform sorted selection by passing 'n' elements over a heap of size 'k'
-          // overall complexity =  O (n * ln(k))
+  int64_t tp_threads = threadpool != nullptr ? threadpool->NumThreads() : 1;
+  int64_t num_threads = std::min(tp_threads, rows);  // split on rows so can't have more threads than rows
 
-          // Build a min-heap/max-heap, the heap element is pair of (value, idx)
-          // The top of the heap is the smallest/largest value depending on whether it is a min-heap/max-heap
-          // This is a min-heap if largest == true, this is a max-heap if largest == false
-          priority_queue<pair<typename Comparator::DataType, int64_t>, vector<pair<typename Comparator::DataType, int64_t>>, Comparator> heap;
+  // rough attempt to make sure there's enough work for each thread. if there's insufficient work the usage of
+  // too many threads degrades performance.
+  // TODO: May want a different calculation for each branch below instead.
+  int64_t threads_needed = static_cast<int64_t>(std::floor(input_shape.Size() * k / (128 * 1024)));
+  num_threads = std::max(std::min(threads_needed, num_threads), static_cast<int64_t>(1));
 
-          // Maintain the size of heap to be less or equal to k, so the
-          // heap will hold the k largest/smallest values
-          for (int64_t l = 0; l < num_blocks; ++l) {
-            const auto value = input_map(i, l * block_slice + j);
-            // largest == true: insert into the min-heap if the size is < k or if the new
-            // element is greater than the min element in the min-heap
+  // from testing various batch sizes relative to k, the following appears to work well as a selector.
+  // tested with following combinations
+  //   batch_size = [ 8, 16, 32, 64, 128, 256, 512, 1024, 2048 ]
+  //            k = [ 1, 2, 4, 6, 8, 16, 24, 32, 48, 64, 128 ]
+  bool use_priority_queue = k != 1 && (k < 4 || (std::log2(k) / std::log2(num_blocks)) < 0.725);
 
-            // largest == false: insert into the min-heap if the size is < k or if the new
-            // element is lesser than the max element in the max-heap
-            if ((heap.size() < k) || (largest && value > heap.top().first) ||
-                (!largest && value < heap.top().first)) {  // the optimizer will clean-up the redundant condition based
-                                                           // on the template parameter 'largest'
-              heap.push({value, l});
-            }
-            if (heap.size() > k) {
-              heap.pop();
+  std::function<void(std::ptrdiff_t batch)> find_top_k;
+
+  if (k == 1) {
+    // just need to compare values and not indexes as the first instance of the best value is always selected
+    find_top_k =
+        [num_threads, rows, block_slice, num_blocks, input_data, cols, &values_map, &indices_map](std::ptrdiff_t batch) {
+          int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
+          int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
+
+          Comparator comparer(input_data);
+
+          for (int64_t i = start_row; i < end_row; ++i) {
+            auto row_offset = i * cols;
+            for (int64_t j = 0; j < block_slice; ++j) {
+              int64_t cur_idx = row_offset + j;
+
+              const auto* cur_value = input_data + cur_idx;  // using pointer to data is faster than input_data[cur_idx]
+              auto best = *cur_value;                        // save best value so we only have one load in the CompareValueOnly call
+              int64_t top_idx = cur_idx;
+
+              for (int64_t l = 1; l < num_blocks; ++l) {
+                cur_value += block_slice;
+                if (comparer.CompareValueOnly(*cur_value, best)) {
+                  best = *cur_value;
+                  top_idx = cur_value - input_data;
+                }
+              }
+
+              values_map(i, j) = best;
+              // convert overall index to result index
+              // avoid '/' if possible for perf reasons
+              indices_map(i, j) = block_slice == 1 ? (top_idx - row_offset - j)
+                                                   : (top_idx - row_offset - j) / block_slice;
             }
           }
-          // Extract these k elements and place them in the results placeholder
-          for (int64_t l = 0; l < k; ++l) {
-            const auto& elem = heap.top();
-            auto col_index = (k - l - 1) * block_slice + j;
-            values_map(i, col_index) = elem.first;
-            indices_map(i, col_index) = elem.second;
-            heap.pop();
+        };
+  } else if (use_priority_queue) {
+    find_top_k =
+        [num_threads, rows, block_slice, num_blocks, k, sorted,
+         input_data, cols, &values_map, &indices_map](std::ptrdiff_t batch) {
+          int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
+          int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
+
+          Comparator comparer(input_data);
+
+          // the heap is stored in indices_data. each iteration overwrites the old data when it adds the
+          // initial k values, so we don't need to clear it.
+          std::vector<int64_t> indices_data(k);
+          int64_t* indices = indices_data.data();  // raw pointer is slightly faster for HeapifyIthPosition
+
+          for (int64_t i = start_row; i < end_row; ++i) {
+            const auto row_offset = i * cols;
+
+            for (int64_t j = 0; j < block_slice; ++j) {
+              int64_t l = 0;
+              auto cur_idx = row_offset + j;
+
+              // add first k items starting from the bottom up
+              for (; l < k; ++l) {
+                indices[k - l - 1] = cur_idx;
+                HeapifyIthPosition(indices, k - l - 1, k, comparer);
+
+                cur_idx += block_slice;
+              }
+
+              // insert remainder if the next value would replace the top of the heap (current worst top k value)
+              // save top so we only have one load in the CompareValueOnly call
+              auto top = input_data[indices[0]];
+              for (; l < num_blocks; ++l) {
+                // we can compare value only. if the current value is equal to the top of the heap it won't
+                // replace it as the index will be higher.
+                if (comparer.CompareValueOnly(input_data[cur_idx], top)) {
+                  indices[0] = cur_idx;
+                  HeapifyIthPosition(indices, 0, k, comparer);
+                  top = input_data[indices[0]];
+                }
+
+                cur_idx += block_slice;
+              }
+
+              if (sorted) {
+                // Extract these k elements and place them in the results placeholder
+                for (l = 0; l < k; ++l) {
+                  auto idx = indices[0];
+                  auto col_index = (k - l - 1) * block_slice + j;
+                  values_map(i, col_index) = input_data[idx];
+                  // convert overall index to result index. avoid '/' if possible for perf reasons
+                  indices_map(i, col_index) = block_slice == 1 ? (idx - row_offset - j)
+                                                               : (idx - row_offset - j) / block_slice;
+
+                  // put the last value at the top of the heap to replace the removed one, and push it into
+                  // place in a heap one smaller.
+                  indices[0] = indices[k - l - 1];
+                  HeapifyIthPosition(indices, 0, k - l - 1, comparer);
+                }
+              } else {
+                for (l = 0; l < k; ++l) {
+                  int64_t idx = indices[l];
+                  auto col_index = l * block_slice + j;
+                  values_map(i, col_index) = input_data[idx];
+                  // convert overall index to result index. avoid '/' if possible for perf reasons
+                  indices_map(i, col_index) = block_slice == 1 ? (idx - row_offset - j)
+                                                               : (idx - row_offset - j) / block_slice;
+                }
+              }
+            }
           }
-        }
-      } else {  // sorted == false
-        // The optimizer will clean-up the redundant condition based on the template parameter 'sorted'
+        };
+  } else {
+    find_top_k =
+        [num_threads, rows, block_slice, num_blocks, k, sorted,
+         input_data, cols,
+         &values_map, &indices_map](std::ptrdiff_t batch) {
+          int64_t start_row = static_cast<int64_t>(batch * rows / num_threads);
+          int64_t end_row = static_cast<int64_t>((batch + 1) * rows / num_threads);
 
-        // If the top K values are not required to be sorted, we use a more optimal selection algorithm
-        // Average - O(n). Worst - O(n * ln(n)) or O(n^2) depending on the implementation, where 'n' is the number of input
+          Comparator comparer(input_data);
 
-        const auto& data_holder = select_top_k<Comparator>(input_map, i, num_blocks, block_slice, j, k, false);
+          // we re-use a single data_holder for performance. avoids allocating memory on each iteration.
+          // the call to SelectTopK overwrites any existing data so we don't need to clear on each iteration.
+          std::vector<int64_t> data_holder(num_blocks);
 
-        // Insert the top 'k' (largest or smallest) elements into the final output buffers
-        for (int64_t l = 0; l < k; ++l) {
-          const auto& elem = data_holder[l];
-          auto col_index = l * block_slice + j;
-          values_map(i, col_index) = elem.first;
-          indices_map(i, col_index) = elem.second;
-        }
-      }
-    }
+          for (int64_t i = start_row; i < end_row; ++i) {
+            auto row_offset = i * cols;
+            for (int64_t j = 0; j < block_slice; ++j) {
+              SelectTopK<Comparator>(comparer, row_offset, num_blocks, block_slice, j, k, sorted, data_holder);
+
+              // Insert the top 'k' (largest or smallest) elements into the final output buffers
+              for (int64_t l = 0; l < k; ++l) {
+                int64_t idx = data_holder[l];
+                auto col_index = l * block_slice + j;
+                values_map(i, col_index) = input_data[idx];
+                // convert overall index to result index. avoid the cost of the '/' is possible
+                indices_map(i, col_index) = block_slice == 1 ? (idx - row_offset - j)
+                                                             : (idx - row_offset - j) / block_slice;
+              }
+            }
+          }
+        };
+  }
+
+  if (num_threads <= 1) {
+    find_top_k(0);
+  } else {
+    // we want to re-use the storage variables in each lambda as much as possible to minimize allocations
+    // on each iteration, so the lambda does multiple rows. e.g. the data_holder and indices_data vectors.
+    // the alternative would be to use TryBatchParallelFor with the lambda doing one row.
+    threadpool->SimpleParallelFor(num_threads, find_top_k);
   }
 }
 
@@ -204,22 +361,14 @@ static Status TopKImpl(OpKernelContext* p_op_kernel_context, const Tensor* input
     return Status::OK();
   }
 
-  if (sorted && largest) {
-    // extract sorted largest TopK elements
-    extract_top_k_elements<true, true, GreaterValueCmp<T>>(input, input_shape, values, indices, output_shape, k,
-                                                           gsl::narrow_cast<unsigned>(axis_parsed));
-  } else if (sorted && !largest) {
-    // extract sorted smallest TopK elements
-    extract_top_k_elements<false, true, LesserValueCmp<T>>(input, input_shape, values, indices, output_shape, k,
-                                                           gsl::narrow_cast<unsigned>(axis_parsed));
-  } else if (largest) {
-    // extract unsorted (order undefined) largest TopK elements
-    extract_top_k_elements<true, false, GreaterValueCmp<T>>(input, input_shape, values, indices, output_shape, k,
-                                                            gsl::narrow_cast<unsigned>(axis_parsed));
+  auto* threadpool = p_op_kernel_context->GetOperatorThreadPool();
+
+  if (largest) {
+    FindTopKElements<GreaterValueCmp<T>>(input, input_shape, values, indices, output_shape, k, sorted,
+                                         gsl::narrow_cast<unsigned>(axis_parsed), threadpool);
   } else {
-    // extract unsorted (order undefined) smallest TopK elements
-    extract_top_k_elements<false, false, LesserValueCmp<T>>(input, input_shape, values, indices, output_shape, k,
-                                                            gsl::narrow_cast<unsigned>(axis_parsed));
+    FindTopKElements<LesserValueCmp<T>>(input, input_shape, values, indices, output_shape, k, sorted,
+                                        gsl::narrow_cast<unsigned>(axis_parsed), threadpool);
   }
 
   return Status::OK();

onnxruntime/core/providers/cpu/rnn/deep_cpu_lstm.cc

@@ -163,7 +163,7 @@ Equations (Default: f=Sigmoid, g=Tanh, h=Tanh):
 namespace onnxruntime {
 
 template <typename TLambda>
-static inline void ExecuteLambdaInParallel(TLambda lambda, int max, int step,
+static inline void ExecuteLambdaInParallel(TLambda lambda, int max, int step, double cost,
                                            onnxruntime::concurrency::ThreadPool* ttp) {
   // #define NOTHREADS to execute the lambdas directly and in order if you need to do that to debug
 
@@ -174,67 +174,12 @@ static inline void ExecuteLambdaInParallel(TLambda lambda, int max, int step,
     std::bind(lambda, i)();
   }
 #else
-
-  // ORT_ENFORCE may and does throw at times from within the tasks that run
-  // on a thread-pool. Without propagating exceptions the process exits silently
-  // which will make diagnosing bugs more difficult.
-
-  // \! UGLY
-  // We have a problem here with the current thread-pool is that it takes std::function
-  // by value and copies it more than once (even though it is movable).
-  //
-  // To report status and exceptions properly it's better to use
-  // futures and promises but they are not copyable, so we can't come up with a functor
-  // with a promise member and we are downgrading to C++11 where we can't have captures that moved in.
-  //
-  // At the same time promises MUST live in the child thread so if we throw from the main thread
-  // we don't destroy any promises that are on the main thread stack which children threads may still be using.
-  //
-  // The only solution with the current Eigen that comes to mind is to have shared_ptr to with std::promise.
-  //
   const int total_tasks = max / (step > 0 ? step : 1) + (max % step > 0 ? 1 : 0);
-  std::vector<std::future<void> > futures;
-  futures.reserve(total_tasks);
-
-  if (ttp != nullptr) {
-    for (int i = 0, t = 0; i < max; i += step, ++t) {
-      auto p_ptr = std::make_shared<std::promise<void> >();
-      futures.push_back(p_ptr->get_future());
-      ttp->Schedule([p_ptr, lambda, i]() {
-        try {
-          lambda(i);
-          p_ptr->set_value();
-        } catch (...) {
-          p_ptr->set_exception(std::current_exception());
-        }
-      });
+  concurrency::ThreadPool::TryParallelFor(ttp, total_tasks, cost, [lambda, step](ptrdiff_t first, ptrdiff_t last) {
+    for (int i = static_cast<int>(first), end = static_cast<int>(last); i < end; ++i) {
+      lambda(i * step);
     }
-
-    // We'd like to wait until all of the tasks have finished
-    // even though one or more have already thrown. We will store
-    // the first exception and then will re-throw at the end.
-    std::exception_ptr pending_exception;
-    for (auto& fut : futures) {
-      try {
-        // get() will re-throw any exceptions
-        // the running task may throw
-        fut.get();
-      } catch (...) {
-        if (!pending_exception) {
-          pending_exception = std::current_exception();
-        }
-      }
-    }
-
-    if (pending_exception) {
-      std::rethrow_exception(pending_exception);
-    }
-  } else {
-    for (int i = 0; i < max; i += step) {
-      std::bind(lambda, i)();
-    }
-  }
-
+  });
 #endif
 }
 
@@ -488,8 +433,7 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const {
 
   const size_t last_cell_size_per_direction = batch_size * hidden_size_;
   IAllocatorUniquePtr<T> local_last_cell;
-  gsl::span<T> last_cell = Y_c ? Y_c->MutableDataAsSpan<T>() :
-                                 Allocate(alloc, last_cell_size_per_direction * num_directions_, local_last_cell);
+  gsl::span<T> last_cell = Y_c ? Y_c->MutableDataAsSpan<T>() : Allocate(alloc, last_cell_size_per_direction * num_directions_, local_last_cell);
 
   gsl::span<T> last_cell_1 = last_cell.subspan(0, last_cell_size_per_direction);
 
@@ -501,17 +445,12 @@ Status DeepCpuLstmOp::ComputeImpl(OpKernelContext& context) const {
         recurrent_weights.subspan(hidden_weights_size_per_direction, hidden_weights_size_per_direction);
     gsl::span<const T> bias_2 = bias.empty() ? bias : bias.subspan(bias_size_per_direction, bias_size_per_direction);
     gsl::span<const T> peephole_weights_2 =
-        peephole_weights.empty() ?
-            peephole_weights :
-            peephole_weights.subspan(peephole_weights_size_per_direction, peephole_weights_size_per_direction);
+        peephole_weights.empty() ? peephole_weights : peephole_weights.subspan(peephole_weights_size_per_direction, peephole_weights_size_per_direction);
 
     gsl::span<const T> initial_hidden_2 =
-        initial_hidden.empty() ?
-            initial_hidden :
-            initial_hidden.subspan(initial_hidden_size_per_direction, initial_hidden_size_per_direction);
+        initial_hidden.empty() ? initial_hidden : initial_hidden.subspan(initial_hidden_size_per_direction, initial_hidden_size_per_direction);
     gsl::span<const T> initial_cell_2 =
-        initial_cell.empty() ? initial_cell :
-                               initial_cell.subspan(initial_cell_size_per_direction, initial_cell_size_per_direction);
+        initial_cell.empty() ? initial_cell : initial_cell.subspan(initial_cell_size_per_direction, initial_cell_size_per_direction);
     gsl::span<T> output_2 =
         output.empty() ? output : output.subspan(per_direction_offset, output_size - per_direction_offset);
 
@@ -861,7 +800,7 @@ void UniDirectionalLstm<T>::Compute(const gsl::span<const T>& inputs_arg,
         ComputeGemm(local_fused_hidden_rows, hidden_size_x4, hidden_size_, alpha, previous_state,
                     previous_state_end,                                                  // Ht-1
                     hidden_size_, recurrent_weights.cbegin(), recurrent_weights.cend(),  // R[iofc]
-                    hidden_size_, beta, step_out_IOFC, output_iofc_.end(),  // input contains Xt*(W[iofc]^T)
+                    hidden_size_, beta, step_out_IOFC, output_iofc_.end(),               // input contains Xt*(W[iofc]^T)
                     hidden_size_x4, nullptr);
 
         DumpMatrix("Xt*(W[iofc]^T) + Ht-t*R[iofc]" + row_str, &*step_out_IOFC, local_fused_hidden_rows, hidden_size_x4);
@@ -910,7 +849,8 @@ void UniDirectionalLstm<T>::Compute(const gsl::span<const T>& inputs_arg,
       }
     };
 
-    ExecuteLambdaInParallel(hidden_gemm_and_activations, batch_size_, fused_hidden_rows, mlas_tp_);
+    double cost = max_sequence_length * fused_hidden_rows;  // TODO: approximate cost, needs more tuning.
+    ExecuteLambdaInParallel(hidden_gemm_and_activations, batch_size_, fused_hidden_rows, cost, mlas_tp_);
 
   } else {
     span_T_const_iter previous_state_end = batched_hidden_state_one_step.cend();
@@ -935,7 +875,7 @@ void UniDirectionalLstm<T>::Compute(const gsl::span<const T>& inputs_arg,
       // calculate Xt*(W[iofc]^T) + Ht-t*R[iofc]
       ComputeGemm(batch_size_, hidden_size_x4, hidden_size_, alpha, previous_state, previous_state_end,  // Ht-1
                   hidden_size_, recurrent_weights.cbegin(), recurrent_weights.cend(),                    // R[iofc]
-                  hidden_size_, beta, step_out_IOFC, output_iofc_.end(),  // input contains Xt*(W[iofc]^T)
+                  hidden_size_, beta, step_out_IOFC, output_iofc_.end(),                                 // input contains Xt*(W[iofc]^T)
                   hidden_size_x4, mlas_tp_);
 
       span_T_iter batched_output;
@@ -1011,7 +951,7 @@ void UniDirectionalLstm<T>::Compute(const gsl::span<const T>& inputs_arg,
 
   if (output_sequence && direction_ == Direction::kReverse)
     ReverseSequence<T>(outputs, original_outputs, sequence_lengths, seq_length_, batch_size_, hidden_size_,
-                       num_directions,mlas_tp_);
+                       num_directions, mlas_tp_);
 }
 
 // #define PREVIOUS_BROKEN_VERSION
@@ -1141,7 +1081,7 @@ void UniDirectionalLstm<T>::GateComputations(
 
 template <typename T>
 void UniDirectionalLstm<T>::SetNumThreads() {
-  int threads = mlas_tp_ == nullptr ? 1 : mlas_tp_->NumThreads();
+  int threads = concurrency::ThreadPool::NumThreads(mlas_tp_);
 
   if (threads < 1)
     threads = 1;

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <fstream>
 #include "core/graph/onnx_protobuf.h"
 
 #include "tensorrt_execution_provider.h"
@@ -124,6 +125,11 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv
   if (!fp16_enable_env.empty()) {
     fp16_enable_ = (std::stoi(fp16_enable_env) == 0 ? false : true);
   }
+
+  const std::string dump_subgraphs_env = env_instance.GetEnvironmentVar(tensorrt_env_vars::kDumpSubgraphs);
+  if (!dump_subgraphs_env.empty()) {
+    dump_subgraphs_ = (std::stoi(dump_subgraphs_env) == 0 ? false : true);
+  }
 }
 
 TensorrtExecutionProvider::~TensorrtExecutionProvider() {}
@@ -252,7 +258,7 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
 
   // Find inputs and outputs of the subgraph
   std::unique_ptr<IndexedSubGraph> sub_graph = onnxruntime::make_unique<IndexedSubGraph>();
-  std::unordered_map<const NodeArg *, int> fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
+  std::unordered_map<const NodeArg*, int> fused_inputs, fused_outputs, fused_outputs_to_add, graph_outputs_to_add;
   std::unordered_set<const NodeArg*> erased;
   int input_order = 0;
   int output_order = 0;
@@ -318,7 +324,7 @@ std::unique_ptr<IndexedSubGraph> TensorrtExecutionProvider::GetSubGraph(SubGraph
   fused_outputs.insert(graph_outputs_to_add.begin(), graph_outputs_to_add.end());
 
   // Sort inputs and outputs by the order they were added
-  std::multimap<int, const NodeArg *> inputs, outputs;
+  std::multimap<int, const NodeArg*> inputs, outputs;
   for (auto it = fused_inputs.begin(), end = fused_inputs.end(); it != end; ++it) {
     inputs.insert(std::pair<int, const NodeArg*>(it->second, it->first));
   }
@@ -384,7 +390,7 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect
         std::vector<std::string> subgraph_output_names;
         for (const auto& index : group.first) {
           const auto& node = graph.GetNode(node_index[index]);
-          std::vector<onnxruntime::NodeArg *> inputs, outputs;
+          std::vector<onnxruntime::NodeArg*> inputs, outputs;
           for (auto input : node->InputDefs()) {
             auto& n_input = graph_build.GetOrCreateNodeArg(input->Name(), input->TypeAsProto());
             inputs.push_back(&n_input);
@@ -661,6 +667,12 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     std::string string_buf;
     model_proto.SerializeToString(&string_buf);
 
+    if (dump_subgraphs_) {
+      // Dump the TensorRT subgraph if enabled via ORT_TENSORRT_DUMP_SUBGRAPHS env variable.
+      std::fstream dump(fused_node->Name() + ".onnx", std::ios::out | std::ios::trunc | std::ios::binary);
+      model_proto.SerializeToOstream(&dump);
+    }
+
     // Create TensorRT engine
     TensorrtLogger& trt_logger = GetTensorrtLogger();
     auto trt_builder = unique_pointer<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(trt_logger));
@@ -1009,4 +1021,4 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
 
   return Status::OK();
 }
-}  // namespace onnxruntime
+}  // namespace onnxruntime

onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h

@@ -16,6 +16,7 @@ static const std::string kMaxPartitionIterations = "ORT_TENSORRT_MAX_PARTITION_I
 static const std::string kMinSubgraphSize = "ORT_TENSORRT_MIN_SUBGRAPH_SIZE";
 static const std::string kMaxWorkspaceSize = "ORT_TENSORRT_MAX_WORKSPACE_SIZE";
 static const std::string kFP16Enable = "ORT_TENSORRT_FP16_ENABLE";
+static const std::string kDumpSubgraphs = "ORT_TENSORRT_DUMP_SUBGRAPHS";
 }  // namespace tensorrt_env_vars
 
 class TensorrtLogger : public nvinfer1::ILogger {
@@ -86,6 +87,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
   int max_partition_iterations_ = 1000;
   int min_subgraph_size_ = 1;
   bool fp16_enable_ = false;
+  bool dump_subgraphs_ = false;
 
   struct InferDeleter {
     template <typename T>
@@ -126,7 +128,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
                                         const onnxruntime::GraphViewer& graph, bool* early_termination) const;
 
   void RemoveTensorRTGraphCycles(SubGraphCollection_t& supported_nodes_vector, const onnxruntime::GraphViewer& graph) const;
-  
+
   AllocatorPtr allocator_;
 };
 

onnxruntime/core/session/abi_session_options.cc

@@ -141,7 +141,15 @@ ORT_API_STATUS_IMPL(OrtApis::SetSessionGraphOptimizationLevel, _In_ OrtSessionOp
 }
 
 ORT_API_STATUS_IMPL(OrtApis::SetIntraOpNumThreads, _Inout_ OrtSessionOptions* options, int intra_op_num_threads) {
+#ifdef _OPENMP
+  ORT_UNUSED_PARAMETER(options);
+  ORT_UNUSED_PARAMETER(intra_op_num_threads);
+  LOGS_DEFAULT(WARNING) << "Since openmp is enabled in this build, this API cannot be used to configure"
+                           " intra op num threads. Please use the openmp environment variables to control"
+                           " the number of threads.";
+#else
   options->value.intra_op_param.thread_pool_size = intra_op_num_threads;
+#endif
   return nullptr;
 }
 

onnxruntime/core/session/environment.cc

@@ -62,12 +62,12 @@ Status Environment::Initialize(std::unique_ptr<logging::LoggingManager> logging_
     if (to.name == nullptr) {
       to.name = ORT_TSTR("intra-op");
     }
-    intra_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, nullptr);
+    intra_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP, nullptr);
     to = tp_options->inter_op_thread_pool_params;
     if (to.name == nullptr) {
       to.name = ORT_TSTR("inter-op");
     }
-    inter_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, nullptr);
+    inter_op_thread_pool_ = concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP, nullptr);
   }
 
   try {

onnxruntime/core/session/inference_session.cc

@@ -183,7 +183,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
                              session_options_.execution_mode == ExecutionMode::ORT_SEQUENTIAL &&
                              to.affinity_vec_len == 0;
       thread_pool_ =
-          concurrency::CreateThreadPool(&Env::Default(), to, nullptr);
+          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP, nullptr);
     }
     if (session_options_.execution_mode == ExecutionMode::ORT_PARALLEL) {
       OrtThreadPoolParams to = session_options_.inter_op_param;
@@ -194,7 +194,7 @@ void InferenceSession::ConstructorCommon(const SessionOptions& session_options,
       if (to.name == nullptr)
         to.name = ORT_TSTR("intra-op");
       inter_op_thread_pool_ =
-          concurrency::CreateThreadPool(&Env::Default(), to, nullptr);
+          concurrency::CreateThreadPool(&Env::Default(), to, concurrency::ThreadPoolType::INTER_OP, nullptr);
       if (inter_op_thread_pool_ == nullptr) {
         LOGS(*session_logger_, INFO) << "Failed to create the inter-op thread pool for the parallel executor, setting ExecutionMode to SEQUENTIAL";
         session_options_.execution_mode = ExecutionMode::ORT_SEQUENTIAL;

onnxruntime/core/util/thread_utils.cc

@@ -9,59 +9,51 @@
 
 namespace onnxruntime {
 namespace concurrency {
-static inline std::vector<size_t> GenerateVectorOfN(size_t n) {
-  std::vector<size_t> ret(n);
-  std::iota(ret.begin(), ret.end(), 0);
-  return ret;
-                }
-#ifdef _WIN32
-		// This function doesn't support systems with more than 64 logical processors
-		static std::vector<size_t> GetNumCpuCores() {
-			// Indeed 64 should be enough. However, it's harmless to have a little more.
-			SYSTEM_LOGICAL_PROCESSOR_INFORMATION buffer[256];
-			DWORD returnLength = sizeof(buffer);
-			if (GetLogicalProcessorInformation(buffer, &returnLength) == FALSE) {
-                          return GenerateVectorOfN(std::thread::hardware_concurrency());
-                        }
-                        std::vector<size_t> ret;
-			int count = (int)(returnLength / sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION));
-			for (int i = 0; i != count; ++i) {
-				if (buffer[i].Relationship == RelationProcessorCore) {
-					ret.push_back(buffer[i].ProcessorMask);
-				}
-			}
-			if (ret.empty())
-                          return GenerateVectorOfN(std::thread::hardware_concurrency());
-                        return ret;
-                }
-#else
-		static std::vector<size_t> GetNumCpuCores() {
-			return GenerateVectorOfN(std::thread::hardware_concurrency() / 2);
-		}
-#endif
-		std::unique_ptr<ThreadPool> CreateThreadPool(Env* env, OrtThreadPoolParams options, Eigen::Allocator* allocator) {
-			if (options.thread_pool_size == 1)
-				return nullptr;
-			std::vector<size_t> cpu_list;
-			ThreadOptions to;
-			if (options.affinity_vec_len != 0) {
-				to.affinity.assign(options.affinity_vec, options.affinity_vec + options.affinity_vec_len);
-			}
-			if (options.thread_pool_size <= 0) {  // default
-				cpu_list = GetNumCpuCores();
-				if (cpu_list.empty() || cpu_list.size() == 1)
-					return nullptr;
-				options.thread_pool_size = static_cast<int>(cpu_list.size());
-				if (options.auto_set_affinity)
-					to.affinity = cpu_list;
-			}
+static std::unique_ptr<ThreadPool>
+CreateThreadPoolHelper(Env* env, OrtThreadPoolParams options, Eigen::Allocator* allocator) {
+  if (options.thread_pool_size == 1)
+    return nullptr;
+  std::vector<size_t> cpu_list;
+  ThreadOptions to;
+  if (options.affinity_vec_len != 0) {
+    to.affinity.assign(options.affinity_vec, options.affinity_vec + options.affinity_vec_len);
+  }
+  if (options.thread_pool_size <= 0) {  // default
+    cpu_list = Env::Default().GetThreadAffinityMasks();
+    if (cpu_list.empty() || cpu_list.size() == 1)
+      return nullptr;
+    options.thread_pool_size = static_cast<int>(cpu_list.size());
+    if (options.auto_set_affinity)
+      to.affinity = cpu_list;
+  }
 
-			return onnxruntime::make_unique<ThreadPool>(env, to, options.name, options.thread_pool_size,
-				options.allow_spinning, allocator);
-		}
-                }  // namespace concurrency
+  return onnxruntime::make_unique<ThreadPool>(env, to, options.name, options.thread_pool_size,
+                                              options.allow_spinning, allocator);
+}
+
+std::unique_ptr<ThreadPool>
+CreateThreadPool(Env* env, OrtThreadPoolParams options, ThreadPoolType tpool_type, Eigen::Allocator* allocator) {
+// If openmp is enabled we don't want to create any additional threadpools for sequential execution.
+// However, parallel execution relies on the existence of a separate threadpool. Hence we allow eigen threadpools
+// to be created for parallel execution.
+#ifdef _OPENMP
+  ORT_UNUSED_PARAMETER(env);
+  ORT_UNUSED_PARAMETER(options);
+  ORT_UNUSED_PARAMETER(allocator);
+  if (tpool_type != ThreadPoolType::INTER_OP) {
+    return nullptr;
+  } else {
+    return CreateThreadPoolHelper(env, options, allocator);
+  }
+#else
+  ORT_UNUSED_PARAMETER(tpool_type);
+  return CreateThreadPoolHelper(env, options, allocator);
+#endif
+}
+
+}  // namespace concurrency
 }  // namespace onnxruntime
-namespace OrtApis{
+namespace OrtApis {
 ORT_API_STATUS_IMPL(CreateThreadingOptions, _Outptr_ OrtThreadingOptions** out) {
   *out = new OrtThreadingOptions();
   return nullptr;
@@ -70,4 +62,4 @@ ORT_API_STATUS_IMPL(CreateThreadingOptions, _Outptr_ OrtThreadingOptions** out)
 ORT_API(void, ReleaseThreadingOptions, _Frees_ptr_opt_ OrtThreadingOptions* p) {
   delete p;
 }
-}
+}  // namespace OrtApis

onnxruntime/core/util/thread_utils.h

@@ -7,7 +7,7 @@
 #include <memory>
 #include <string>
 
-struct OrtThreadPoolParams{
+struct OrtThreadPoolParams {
   //0: Use default setting. (All the physical cores or half of the logical cores)
   //1: Don't create thread pool
   //n: Create a thread pool with n threads.
@@ -25,7 +25,7 @@ struct OrtThreadPoolParams{
   size_t* affinity_vec = nullptr;
   size_t affinity_vec_len = 0;
   const ORTCHAR_T* name = nullptr;
-} ;
+};
 
 struct OrtThreadingOptions {
   // Params for creating the threads that parallelizes execution of an op
@@ -33,13 +33,17 @@ struct OrtThreadingOptions {
 
   // Params for creating the threads that parallelizes execution across ops
   OrtThreadPoolParams inter_op_thread_pool_params;
-} ;
+};
 
 namespace onnxruntime {
 
 namespace concurrency {
-
+enum class ThreadPoolType : uint8_t {
+  INTRA_OP,
+  INTER_OP
+};
 std::unique_ptr<ThreadPool> CreateThreadPool(Env* env, OrtThreadPoolParams options,
+                                             ThreadPoolType tpool_type,
                                              Eigen::Allocator* allocator = nullptr);
 }  // namespace concurrency
 }  // namespace onnxruntime

onnxruntime/python/tools/bert/BertOnnxModel.py

@@ -30,20 +30,14 @@ class BertOptimizationOptions:
 
 class BertOnnxModel(OnnxModel):
 
-    def __init__(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
+    def __init__(self, model, num_heads, hidden_size):
         assert num_heads > 0
         assert hidden_size % num_heads == 0
-        assert sequence_length > 0
 
         super(BertOnnxModel, self).__init__(model)
         self.num_heads = num_heads
-        self.sequence_length = sequence_length
         self.hidden_size = hidden_size
 
-        self.input_int32 = input_int32
-        self.gpu_only = gpu_only
-        self.float16 = float16
-
         # A lookup table with mask input as key, and mask index output as value
         self.mask_indice = {}
         # A lookup table with mask input as key, and cast (to int32) output as value
@@ -72,13 +66,13 @@ class BertOnnxModel(OnnxModel):
         graph_input = self.find_graph_input(input_name)
         if graph_input is not None and graph_input.type.tensor_type.elem_type != TensorProto.INT32:
             cast_output, cast_node = self.cast_input_to_int32(input_name)
-            logger.debug("Casted graph input {input_name} to int32")
+            logger.debug(f"Casted graph input {input_name} to int32")
             return True, cast_output
 
         logger.debug(f"Did not cast graph input {input_name} to int32: found {graph_input is not None}")
         return False, input_name
 
-    def undo_cast_input_to_int32(self, input_name):
+    def remove_cast_int32(self, input_name):
         input_name_to_nodes = self.input_name_to_nodes()
         nodes = input_name_to_nodes[input_name]
         for node in nodes:
@@ -830,14 +824,12 @@ class BertOnnxModel(OnnxModel):
 
         # Cast input_ids and segment_ids to int32.
         if self.find_graph_input(input_ids):
-            if not self.input_int32:
-                casted, input_ids = self.cast_graph_input_to_int32(input_ids)
+            casted, input_ids = self.cast_graph_input_to_int32(input_ids)
         else:
             input_ids, input_ids_cast_node = self.cast_input_to_int32(input_ids)
 
         if self.find_graph_input(segment_ids):
-            if not self.input_int32:
-                casted, segment_ids = self.cast_graph_input_to_int32(segment_ids)
+            casted, segment_ids = self.cast_graph_input_to_int32(segment_ids)
         else:
             segment_ids, segment_ids_cast_node = self.cast_input_to_int32(segment_ids)
 
@@ -874,8 +866,6 @@ class BertOnnxModel(OnnxModel):
         self.replace_input_of_all_nodes(normalize_node.output[0], 'embed_output')
 
         self.remove_nodes(nodes_to_remove)
-        self.add_node(embed_node)
-        self.prune_graph()
 
         return embed_node
 
@@ -887,67 +877,71 @@ class BertOnnxModel(OnnxModel):
 
         if len(self.mask_indice) > 1:
             logger.info("There are multiple mask inputs found!")
-
-        if len(self.mask_indice) != 1:
+        elif len(self.mask_indice) != 1:
             logger.info("Fused EmbedLayerNormalization (no mask) count: 1")
-            return
+        else:
+            mask_input_name = next(iter(self.mask_indice))
+            mask_output_name = self.mask_indice[mask_input_name]
+            output_name_to_node = self.output_name_to_node()
+            mask_node = output_name_to_node[mask_output_name]
 
-        mask_input_name = next(iter(self.mask_indice))
-        mask_output_name = self.mask_indice[mask_input_name]
-        mask_node = output_name_to_node[mask_output_name]
+            nodes_to_remove = []
+            nodes_to_remove.extend([mask_node])
 
-        nodes_to_remove = []
-        nodes_to_remove.extend([mask_node])
+            # store inputs for further processing
+            self.bert_inputs.append(mask_input_name)
 
-        # store inputs for further processing
-        self.bert_inputs.append(mask_input_name)
-
-        if not self.input_int32:
             # When mask has been casted to int32, use that casted one as input of embed layer norm.
             if mask_input_name in self.mask_casted:
                 mask_input_name = self.mask_casted[mask_input_name]
-        else:
-            self.undo_cast_input_to_int32(mask_input_name)
 
-        embed_node.input[7] = mask_input_name
-        embed_node.output[1] = mask_output_name
-        logger.info("Added mask to EmbedLayerNormalization")
+            embed_node.input.append(mask_input_name)
+            embed_node.output[1] = mask_output_name
+            logger.info("Added mask to EmbedLayerNormalization")
+            logger.info("Fused EmbedLayerNormalization count: 1")
 
-        # Change graph input data type int32 if needed.
-        if self.input_int32:
-            self.change_input_to_int32()
-
-        logger.info("Fused EmbedLayerNormalization count: 1")
+        self.add_node(embed_node)
+        self.prune_graph()
 
     def get_bert_inputs(self, include_mask=True):
         return self.bert_inputs if include_mask else self.bert_inputs[:2]
 
-    def get_batch_size_from_graph_input(self):
+    def get_bert_input_shape(self):
         graph = self.graph()
         bert_inputs = self.get_bert_inputs()
         for input in graph.input:
             if input.name in bert_inputs:
                 tensor_type = input.type.tensor_type
                 if (tensor_type.HasField("shape")):
-                    for d in tensor_type.shape.dim:
-                        if (d.HasField("dim_value")):
-                            return d.dim_value
-                        elif (d.HasField("dim_param")):
-                            return str(d.dim_param)  # unknown dimension with symbolic name
-                        return None
-        return None
+                    batch_size = None
+                    d = tensor_type.shape.dim[0]
+                    if (d.HasField("dim_value")):
+                        batch_size = d.dim_value
+                    elif (d.HasField("dim_param")):
+                        batch_size =  str(d.dim_param)
+
+                    sequence_length = None
+                    d = tensor_type.shape.dim[1]
+                    if (d.HasField("dim_value")):
+                        sequence_length = d.dim_value
+                    elif (d.HasField("dim_param")):
+                        sequence_length = str(d.dim_param)
+                    return batch_size, sequence_length
+
+        return None, None
 
     def change_input_to_int32(self):
         original_opset_version = self.model.opset_import[0].version
         graph = self.graph()
 
-        batch_size = self.get_batch_size_from_graph_input()
+        batch_size, sequence_length = self.get_bert_input_shape()
         new_graph_inputs = []
 
         bert_inputs = self.get_bert_inputs()
         for input in graph.input:
             if input.name in bert_inputs:
-                input_shape = [batch_size if isinstance(batch_size, int) else 1, self.sequence_length]
+                self.remove_cast_int32(input.name)
+                input_shape = [batch_size if isinstance(batch_size, int) else 1, sequence_length if isinstance(sequence_length, int) else 128]
                 int32_input = onnx.helper.make_tensor_value_info(input.name, TensorProto.INT32, input_shape)
                 new_graph_inputs.append(int32_input)
             else:
@@ -962,8 +956,8 @@ class BertOnnxModel(OnnxModel):
 
         self.model = onnx.helper.make_model(graph_def, producer_name='bert model optimizer')
 
-        if isinstance(batch_size, str):
-            self.use_dynamic_axes(batch_size, None)
+        if isinstance(batch_size, str) or isinstance(sequence_length, str):
+            self.use_dynamic_axes(batch_size if isinstance(batch_size, str) else None, sequence_length if isinstance(sequence_length, str) else None)
 
         # restore opset version
         self.model.opset_import[0].version = original_opset_version
@@ -1165,9 +1159,6 @@ class BertOnnxModel(OnnxModel):
             # Fuse SkipLayerNormalization and Add Bias before it.
             self.fuse_add_bias_skip_layer_norm()
 
-        if self.float16:
-            self.convert_model_float32_to_float16()
-
         self.remove_unused_constant()
 
         # Use symbolic batch dimension in input and output.

onnxruntime/python/tools/bert/BertOnnxModelKeras.py

@@ -17,8 +17,8 @@ logger = logging.getLogger(__name__)
 
 class BertOnnxModelKeras(BertOnnxModelTF):
 
-    def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
-        super().__init__(model, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
+    def __init(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
 
     def match_mask_path(self, add_or_sub_before_softmax):
         mask_nodes = self.match_parent_path(add_or_sub_before_softmax, ['Mul', 'Sub', 'Reshape', 'Cast'],

onnxruntime/python/tools/bert/BertOnnxModelTF.py

@@ -17,8 +17,8 @@ logger = logging.getLogger(__name__)
 
 class BertOnnxModelTF(BertOnnxModel):
 
-    def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
-        super().__init__(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
+    def __init(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
 
     """
      Fuse Gelu with Erf into one node:

onnxruntime/python/tools/bert/Gpt2OnnxModel.py

@@ -16,8 +16,8 @@ logger = logging.getLogger(__name__)
 
 class Gpt2OnnxModel(BertOnnxModel):
 
-    def __init(self, model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only):
-        super().__init__(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
+    def __init(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
 
     def fuse_attention(self):
         """

onnxruntime/python/tools/bert/README.md

@@ -45,27 +45,23 @@ For tf2onnx, please refer to this notebook: https://github.com/onnx/tensorflow-o
 
 Example of using the script bert_model_optimization.py to convert a BERT-large model to run in V100 GPU:
 ```console
-python bert_model_optimization.py --input original_model.onnx --output optimized_model_gpu.onnx --num_heads 16 --hidden_size 1024 --input_int32 --float16 --gpu_only
+python bert_model_optimization.py --input original_model.onnx --output optimized_model_gpu.onnx --num_heads 16 --hidden_size 1024 --input_int32 --float16
 ```
 
 ### Options
 
-See below for description of all the options of bert_model_optimization.py:
+See below for description of some options of bert_model_optimization.py:
 
 - **input**: input model path
 - **output**: output model path
 - **model_type**: (*defaul: bert*)
-    There are 3 model types: *bert*, *bert_tf* and *bert_keras* for models exported by PyTorch, tf2onnx and keras2onnx respectively.
+    There are 4 model types: *bert* (exported by PyTorch), *bert_tf* (BERT exported by PyTorch), *bert_keras* (BERT exported by keras2onnx) and *gpt2* (exported by PyTorch) respectively.
 - **num_heads**: (*default: 12*)
     Number of attention heads. BERT-base and BERT-large has 12 and 16 respectively.
 - **hidden_size**: (*default: 768*)
     BERT-base and BERT-large has 768 and 1024 hidden nodes respectively.
-- **sequence_length**: (*default: 128*)
-    Maximum sequence length.
 - **input_int32**: (*optional*)
     Exported model ususally uses int64 tensor as input. If this flag is specified, int32 tensors will be used as input, and it could avoid un-necessary Cast nodes and get better performance.
-- **gpu_only**: (*optional*)
-    Specify the option if running on GPU only.
 - **float16**: (*optional*)
     By default, model uses float32 in computation. If this flag is specified, half-precision float will be used. This option is recommended for NVidia GPU with Tensor Core like V100 and T4. For older GPUs, float32 is likely faster.
 - **verbose**: (*optional*)
@@ -77,12 +73,13 @@ Right now, this tool assumes input model has 3 inputs for input IDs, segment IDs
 
 Most optimizations require exact match of a subgraph. That means this tool could only support similar models with such subgraphs. Any layout change in subgraph might cause optimization not working. Note that different training or export tool (including different versions) might get different graph layouts.
 
-Here is list of models that have been tested using this tool:
+Here is list of models from [Huggingface Transformers](https://github.com/huggingface/transformers/) that have been tested using this tool:
 - **BertForSequenceClassification** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_glue.py) exported by PyTorch 1.2-1.4 using opset version 10 or 11.
 - **BertForQuestionAnswering** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_squad.py) exported by PyTorch 1.2-1.4 using opset version 10 or 11.
 - **TFBertForSequenceClassification** as in [transformers example](https://github.com/huggingface/transformers/blob/master/examples/run_tf_glue.py) exported by keras2onnx installed from its master source.
-- **TFBertForQuestionAnswering** as in [transformers](https://github.com/huggingface/transformers/) exported by keras2onnx installed from its master source.
-
+- **TFBertForQuestionAnswering** exported by keras2onnx installed from its master source.
+- **GPT2Model** exported by PyTorch 1.4 using opset version 10 or 11.
+- **GPT2LMHeadModel** exported by PyTorch 1.4 using opset version 10 or 11.
 If your model is not in the list, the optimized model might not work. You are welcome to update the scripts to support new models.
 
 ## Model Verification
@@ -102,11 +99,9 @@ pip install onnxruntime-gpu
 python compare_bert_results.py --baseline_model original_model.onnx --optimized_model optimized_model_gpu.onnx --batch_size 1 --sequence_length 128 --samples 100 --use_gpu
 ```
 
-To use onnxruntime-gpu 1.1.*, it is required to install CUDA and cuDNN and add their bin directories to PATH environment variable.
+To use onnxruntime-gpu, it is required to install CUDA and cuDNN and add their bin directories to PATH environment variable.
 
-## Performance Test
-
-The script for model verification will create a sub-directory like batch_1_seq_128 on the directory of optimized model. You can copy the original or optimized model to the sub-directory, and use onnxruntime_perf_test.exe to test performance of C API.
+## Performance Test (Python)
 
 bert_perf_test.py can be used to check the model inference performance of python API. Below are examples:
 
@@ -121,3 +116,19 @@ python bert_perf_test.py --model optimized_model_gpu.onnx --batch_size 1 --seque
 
 After test is finished, a file like perf_results_CPU_B1_S128_<date_time>.txt or perf_results_GPU_B1_S128_<date_time>.txt will be output to the model directory.
 
+## Performance Test (C API)
+
+First, we need generate some test data. Please make sure there is no sub-directories on the directory of onnx model.
+
+Here is an example:
+```console
+python bert_test_data.py --model bert.onnx --batch_size 1 --sequence_length 32 --samples 100 --output_dir .
+```
+
+You can go to root of this git repository, and build onnxruntime_perf_test.exe from source to test performance of C API. Example commands in Windows:
+```console
+build.bat --config RelWithDebInfo --enable_lto --use_openmp --build_shared_lib --parallel --cmake_generator "Visual Studio 16 2019"
+Set OMP_NUM_THREADS=%NUMBER_OF_PROCESSORS%
+Set OMP_WAIT_POLICY=PASSIVE
+build\Windows\RelWithDebInfo\RelWithDebInfo\onnxruntime_perf_test.exe -e cpu -r 100 -s -o 2 bert.onnx output.txt
+```

onnxruntime/python/tools/bert/benchmark_gpt2.py

@@ -142,7 +142,7 @@ def setup_logger(verbose=True):
     logger.setLevel(logging_level)
 
 
-def remove_past_outputs(export_model_path):
+def remove_past_outputs(export_model_path, output_model_path):
     from onnx import ModelProto
     from OnnxModel import OnnxModel
 
@@ -155,9 +155,8 @@ def remove_past_outputs(export_model_path):
     keep_output_names = [bert_model.model.graph.output[0].name]
     logger.info(f"Prune graph to keep the first output and drop past state outputs:{keep_output_names}")
     bert_model.prune_graph(keep_output_names)
-    onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_out1.onnx'.format(int(enable_past_input)))
-    bert_model.save_model_to_file(onnx_model_path)
-    return onnx_model_path
+
+    bert_model.save_model_to_file(output_model_path)
 
 
 def main():
@@ -226,19 +225,20 @@ def main():
     setup_environment(args.use_openmp)
     import onnxruntime
 
-    onnx_model_path = export_model_path if enable_past_input else remove_past_outputs(export_model_path)
+    if enable_past_input:
+        onnx_model_path = export_model_path
+    else:
+        onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_out1.onnx'.format(int(enable_past_input)))
+        remove_past_outputs(export_model_path, onnx_model_path)
 
     if args.enable_optimization:
         from bert_model_optimization import optimize_model
         m = optimize_model(onnx_model_path,
                            model_type='gpt2',
-                           gpu_only=False,
                            num_heads=12,
                            hidden_size=768,
-                           sequence_length=64,
-                           input_int32=False,
-                           float16=False,
-                           opt_level=0)
+                           opt_level=0,
+                           optimization_options=None)
         onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_optimized.onnx'.format(int(enable_past_input)))
         m.save_model_to_file(onnx_model_path)
 

onnxruntime/python/tools/bert/bert_model_optimization.py

@@ -47,7 +47,7 @@ MODEL_CLASSES = {
 }
 
 
-def optimize_by_onnxruntime(onnx_model_path, use_gpu, optimized_model_path=None, opt_level=99):
+def optimize_by_onnxruntime(onnx_model_path, use_gpu=False, optimized_model_path=None, opt_level=99):
     """
     Use onnxruntime package to optimize model. It could support models exported by PyTorch.
 
@@ -116,8 +116,6 @@ def parse_arguments():
                         default=768,
                         help="bert model hidden size. 768 for bert-base model and 1024 for bert-large")
 
-    parser.add_argument('--sequence_length', required=False, type=int, default=128, help="max sequence length")
-
     parser.add_argument('--input_int32',
                         required=False,
                         action='store_true',
@@ -131,12 +129,6 @@ def parse_arguments():
         help="If your target device is V100 or T4 GPU, use this to convert float32 to float16 for best performance")
     parser.set_defaults(float16=False)
 
-    parser.add_argument('--gpu_only',
-                        required=False,
-                        action='store_true',
-                        help="whether the target device is gpu or not")
-    parser.set_defaults(gpu_only=False)
-
     parser.add_argument('--disable_attention', required=False, action='store_true', help="disable Attention fusion")
     parser.set_defaults(disable_attention=False)
 
@@ -196,19 +188,15 @@ def get_optimization_options(args):
 
 def optimize_model(input,
                    model_type,
-                   gpu_only,
                    num_heads,
                    hidden_size,
-                   sequence_length,
-                   input_int32,
-                   float16,
                    opt_level=99,
                    optimization_options=None):
     (optimizer_class, producer, run_onnxruntime) = MODEL_CLASSES[model_type]
 
     input_model_path = input
     if run_onnxruntime and opt_level > 0:
-        input_model_path = optimize_by_onnxruntime(input_model_path, gpu_only, opt_level=opt_level)
+        input_model_path = optimize_by_onnxruntime(input_model_path, use_gpu=False, opt_level=opt_level)
         logger.info("Use OnnxRuntime to optimize and save the optimized model to {}".format(input_model_path))
 
     model = ModelProto()
@@ -223,7 +211,7 @@ def optimize_model(input,
     if optimization_options is None:
         optimization_options = BertOptimizationOptions(model_type)
 
-    bert_model = optimizer_class(model, num_heads, hidden_size, sequence_length, input_int32, float16, gpu_only)
+    bert_model = optimizer_class(model, num_heads, hidden_size)
     bert_model.optimize(optimization_options)
 
     return bert_model
@@ -243,9 +231,14 @@ def main():
 
     optimization_options = get_optimization_options(args)
 
-    bert_model = optimize_model(args.input, args.model_type, args.gpu_only, args.num_heads, args.hidden_size,
-                                args.sequence_length, args.input_int32, args.float16, args.opt_level,
-                                optimization_options)
+    bert_model = optimize_model(args.input, args.model_type, args.num_heads, args.hidden_size,
+                                args.opt_level, optimization_options)
+
+    if args.float16:
+        bert_model.convert_model_float32_to_float16()
+
+    if args.input_int32:
+        bert_model.change_input_to_int32()
 
     bert_model.save_model_to_file(args.output)
 

onnxruntime/python/tools/bert/notebooks/Inference_GPT2_with_OnnxRuntime_on_CPU.ipynb

@@ -14,7 +14,9 @@
    "source": [
     "# Inference PyTorch GPT2 Model with ONNX Runtime on CPU\n",
     "\n",
-    "In this tutorial, you'll be introduced to how to load a GPT2 model from PyTorch, convert it to ONNX, and inference it using ONNX Runtime."
+    "In this tutorial, you'll be introduced to how to load a GPT2 model from PyTorch, convert it to ONNX, and inference it using ONNX Runtime.\n",
+    "\n",
+    "**Note: this work is still in progresss. Need install ort_nightly package before onnxruntime 1.3.0 is ready. The performance number of ort_nightly does not reflect the final result for onnxruntime 1.3.0. **"
    ]
   },
   {
@@ -44,7 +46,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -54,7 +56,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -84,114 +86,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "   benchmark_gpt2.py: no environment variable of OMP_NUM_THREADS\n",
-      "   benchmark_gpt2.py: no environment variable of OMP_WAIT_POLICY\n",
-      "tokenization_utils.py: loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json from cache at ./gpt2\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71\n",
-      "tokenization_utils.py: loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt from cache at ./gpt2\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda\n",
-      "configuration_utils.py: loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json from cache at ./gpt2\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.699bbd1c449e9861456f359d6daa51bd523ac085b4b531ab0aad5a55d091e942\n",
-      "configuration_utils.py: Model config GPT2Config {\n",
-      "  \"architectures\": [\n",
-      "    \"GPT2LMHeadModel\"\n",
-      "  ],\n",
-      "  \"attn_pdrop\": 0.1,\n",
-      "  \"bos_token_id\": null,\n",
-      "  \"do_sample\": false,\n",
-      "  \"embd_pdrop\": 0.1,\n",
-      "  \"eos_token_ids\": null,\n",
-      "  \"finetuning_task\": null,\n",
-      "  \"id2label\": {\n",
-      "    \"0\": \"LABEL_0\",\n",
-      "    \"1\": \"LABEL_1\"\n",
-      "  },\n",
-      "  \"initializer_range\": 0.02,\n",
-      "  \"is_decoder\": false,\n",
-      "  \"label2id\": {\n",
-      "    \"LABEL_0\": 0,\n",
-      "    \"LABEL_1\": 1\n",
-      "  },\n",
-      "  \"layer_norm_epsilon\": 1e-05,\n",
-      "  \"length_penalty\": 1.0,\n",
-      "  \"max_length\": 20,\n",
-      "  \"model_type\": \"gpt2\",\n",
-      "  \"n_ctx\": 1024,\n",
-      "  \"n_embd\": 768,\n",
-      "  \"n_head\": 12,\n",
-      "  \"n_layer\": 12,\n",
-      "  \"n_positions\": 1024,\n",
-      "  \"num_beams\": 1,\n",
-      "  \"num_labels\": 2,\n",
-      "  \"num_return_sequences\": 1,\n",
-      "  \"output_attentions\": false,\n",
-      "  \"output_hidden_states\": false,\n",
-      "  \"output_past\": true,\n",
-      "  \"pad_token_id\": null,\n",
-      "  \"pruned_heads\": {},\n",
-      "  \"repetition_penalty\": 1.0,\n",
-      "  \"resid_pdrop\": 0.1,\n",
-      "  \"summary_activation\": null,\n",
-      "  \"summary_first_dropout\": 0.1,\n",
-      "  \"summary_proj_to_labels\": true,\n",
-      "  \"summary_type\": \"cls_index\",\n",
-      "  \"summary_use_proj\": true,\n",
-      "  \"temperature\": 1.0,\n",
-      "  \"top_k\": 50,\n",
-      "  \"top_p\": 1.0,\n",
-      "  \"torchscript\": false,\n",
-      "  \"use_bfloat16\": false,\n",
-      "  \"vocab_size\": 50257\n",
-      "}\n",
-      "\n",
-      "   modeling_utils.py: loading weights file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin from cache at ./gpt2\\4295d67f022061768f4adc386234dbdb781c814c39662dd1662221c309962c55.778cf36f5c4e5d94c8cd9cefcf2a580c8643570eb327f0d4a1f007fab2acbdf1\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "D:\\Anaconda3\\envs\\cpu_env\\lib\\site-packages\\transformers\\modeling_gpt2.py:143: TracerWarning: Converting a tensor to a Python float might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
-      "  w = w / math.sqrt(v.size(-1))\n",
-      "D:\\Anaconda3\\envs\\cpu_env\\lib\\site-packages\\transformers\\modeling_gpt2.py:145: TracerWarning: Converting a tensor to a Python index might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!\n",
-      "  b = self.bias[:, :, ns - nd : ns, :ns]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "   benchmark_gpt2.py: PyTorch Inference time = 36.77 ms\n",
-      "   benchmark_gpt2.py: OMP_NUM_THREADS=1\n",
-      "   benchmark_gpt2.py: OMP_WAIT_POLICY=ACTIVE\n",
-      "    BertOnnxModel.py: Fused LayerNormalization count: 25\n",
-      "    BertOnnxModel.py: Fused Gelu (FastGelu fits better) count: 12\n",
-      "    BertOnnxModel.py: Fused Reshape count:48\n",
-      "    BertOnnxModel.py: Fused SkipLayerNormalization count: 1\n",
-      "        OnnxModel.py: Removed unused constant nodes: 422\n",
-      "    BertOnnxModel.py: Fused Attention count:0\n",
-      "    BertOnnxModel.py: skip embed layer fusion since mask input is not found\n",
-      "    BertOnnxModel.py: opset verion: 11\n",
-      "        OnnxModel.py: Output model to ./gpt2_onnx\\gpt2_past0_optimized.onnx\n",
-      "   benchmark_gpt2.py: session option: intra_op_num_threads=12\n",
-      "   benchmark_gpt2.py: Start inferencing onnx model: ./gpt2_onnx\\gpt2_past0_optimized.onnx\n",
-      "   benchmark_gpt2.py: OnnxRuntime Inference time = 31.60 ms\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Assume you have git clone the repository of onnxruntime from github.\n",
     "bert_tools_dir = r'D:\\Git\\onnxruntime\\onnxruntime\\python\\tools\\bert'\n",
     "benchmark_script = os.path.join(bert_tools_dir, 'benchmark_gpt2.py')\n",
     "\n",
     "if enable_past_input:\n",
-    "    %run $benchmark_script --cache_dir $cache_dir --output_dir $output_dir --enable_optimization --enable_past_input\n",
+    "    %run $benchmark_script --model_type gpt2 --cache_dir $cache_dir --output_dir $output_dir --enable_optimization --enable_past_input\n",
     "else:\n",
-    "    %run $benchmark_script --cache_dir $cache_dir --output_dir $output_dir --enable_optimization"
+    "    %run $benchmark_script --model_type gpt2 --cache_dir $cache_dir --output_dir $output_dir --enable_optimization"
    ]
   },
   {
@@ -207,269 +113,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "tokenization_utils.py: loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-vocab.json from cache at ./gpt2\\f2808208f9bec2320371a9f5f891c184ae0b674ef866b79c58177067d15732dd.1512018be4ba4e8726e41b9145129dc30651ea4fec86aa61f4b9f40bf94eac71\n",
-      "tokenization_utils.py: loading file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-merges.txt from cache at ./gpt2\\d629f792e430b3c76a1291bb2766b0a047e36fae0588f9dbc1ae51decdff691b.70bec105b4158ed9a1747fea67a43f5dee97855c64d62b6ec3742f4cfdb5feda\n",
-      "configuration_utils.py: loading configuration file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json from cache at ./gpt2\\4be02c5697d91738003fb1685c9872f284166aa32e061576bbe6aaeb95649fcf.699bbd1c449e9861456f359d6daa51bd523ac085b4b531ab0aad5a55d091e942\n",
-      "configuration_utils.py: Model config GPT2Config {\n",
-      "  \"architectures\": [\n",
-      "    \"GPT2LMHeadModel\"\n",
-      "  ],\n",
-      "  \"attn_pdrop\": 0.1,\n",
-      "  \"bos_token_id\": null,\n",
-      "  \"do_sample\": false,\n",
-      "  \"embd_pdrop\": 0.1,\n",
-      "  \"eos_token_ids\": null,\n",
-      "  \"finetuning_task\": null,\n",
-      "  \"id2label\": {\n",
-      "    \"0\": \"LABEL_0\",\n",
-      "    \"1\": \"LABEL_1\"\n",
-      "  },\n",
-      "  \"initializer_range\": 0.02,\n",
-      "  \"is_decoder\": false,\n",
-      "  \"label2id\": {\n",
-      "    \"LABEL_0\": 0,\n",
-      "    \"LABEL_1\": 1\n",
-      "  },\n",
-      "  \"layer_norm_epsilon\": 1e-05,\n",
-      "  \"length_penalty\": 1.0,\n",
-      "  \"max_length\": 20,\n",
-      "  \"model_type\": \"gpt2\",\n",
-      "  \"n_ctx\": 1024,\n",
-      "  \"n_embd\": 768,\n",
-      "  \"n_head\": 12,\n",
-      "  \"n_layer\": 12,\n",
-      "  \"n_positions\": 1024,\n",
-      "  \"num_beams\": 1,\n",
-      "  \"num_labels\": 2,\n",
-      "  \"num_return_sequences\": 1,\n",
-      "  \"output_attentions\": false,\n",
-      "  \"output_hidden_states\": false,\n",
-      "  \"output_past\": true,\n",
-      "  \"pad_token_id\": null,\n",
-      "  \"pruned_heads\": {},\n",
-      "  \"repetition_penalty\": 1.0,\n",
-      "  \"resid_pdrop\": 0.1,\n",
-      "  \"summary_activation\": null,\n",
-      "  \"summary_first_dropout\": 0.1,\n",
-      "  \"summary_proj_to_labels\": true,\n",
-      "  \"summary_type\": \"cls_index\",\n",
-      "  \"summary_use_proj\": true,\n",
-      "  \"temperature\": 1.0,\n",
-      "  \"top_k\": 50,\n",
-      "  \"top_p\": 1.0,\n",
-      "  \"torchscript\": false,\n",
-      "  \"use_bfloat16\": false,\n",
-      "  \"vocab_size\": 50257\n",
-      "}\n",
-      "\n",
-      "   modeling_utils.py: loading weights file https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin from cache at ./gpt2\\4295d67f022061768f4adc386234dbdb781c814c39662dd1662221c309962c55.778cf36f5c4e5d94c8cd9cefcf2a580c8643570eb327f0d4a1f007fab2acbdf1\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "GPT2Model(\n",
-       "  (wte): Embedding(50257, 768)\n",
-       "  (wpe): Embedding(1024, 768)\n",
-       "  (drop): Dropout(p=0.1, inplace=False)\n",
-       "  (h): ModuleList(\n",
-       "    (0): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (1): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (2): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (3): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (4): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (5): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (6): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (7): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (8): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (9): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (10): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "    (11): Block(\n",
-       "      (ln_1): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (attn): Attention(\n",
-       "        (c_attn): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (attn_dropout): Dropout(p=0.1, inplace=False)\n",
-       "        (resid_dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "      (ln_2): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       "      (mlp): MLP(\n",
-       "        (c_fc): Conv1D()\n",
-       "        (c_proj): Conv1D()\n",
-       "        (dropout): Dropout(p=0.1, inplace=False)\n",
-       "      )\n",
-       "    )\n",
-       "  )\n",
-       "  (ln_f): LayerNorm((768,), eps=1e-05, elementwise_affine=True)\n",
-       ")"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "from transformers import GPT2Model, GPT2Tokenizer\n",
     "model_class, tokenizer_class,  model_name_or_path = (GPT2Model,  GPT2Tokenizer,  'gpt2')\n",
@@ -480,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -527,14 +173,15 @@
     "    ort_outputs = onnxruntime_inference(ort_session, input_ids, past, total_runs)\n",
     "    if verify_outputs:\n",
     "        print('PyTorch and OnnxRuntime output 0 (last_state) are close:'.format(0), numpy.allclose(ort_outputs[0], outputs[0].cpu(), rtol=1e-05, atol=1e-04))\n",
-    "        \n",
-    "        for layer in range(model.config.n_layer):\n",
-    "            print('PyTorch and OnnxRuntime layer {} state (present_{}) are close:'.format(layer, layer), numpy.allclose(ort_outputs[1 + layer], outputs[1][layer].cpu(), rtol=1e-05, atol=1e-04))    "
+    "\n",
+    "        if enable_past_input:\n",
+    "            for layer in range(model.config.n_layer):\n",
+    "                print('PyTorch and OnnxRuntime layer {} state (present_{}) are close:'.format(layer, layer), numpy.allclose(ort_outputs[1 + layer], outputs[1][layer].cpu(), rtol=1e-05, atol=1e-04))    "
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -550,7 +197,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -590,6 +237,35 @@
     "                  verbose=False)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def remove_past_outputs(export_model_path, output_model_path):\n",
+    "    from onnx import ModelProto\n",
+    "    from OnnxModel import OnnxModel\n",
+    "\n",
+    "    model = ModelProto()\n",
+    "    with open(export_model_path, \"rb\") as f:\n",
+    "        model.ParseFromString(f.read())\n",
+    "    bert_model = OnnxModel(model)\n",
+    "\n",
+    "    # remove past state outputs and only keep the first output.\n",
+    "    keep_output_names = [bert_model.model.graph.output[0].name]\n",
+    "    logger.info(f\"Prune graph to keep the first output and drop past state outputs:{keep_output_names}\")\n",
+    "    bert_model.prune_graph(keep_output_names)\n",
+    "\n",
+    "    bert_model.save_model_to_file(output_model_path)\n",
+    "    \n",
+    "if enable_past_input:\n",
+    "    onnx_model_path = export_model_path\n",
+    "else:\n",
+    "    onnx_model_path = os.path.join(output_dir, 'gpt2_past{}_out1.onnx'.format(int(enable_past_input)))\n",
+    "    remove_past_outputs(export_model_path, onnx_model_path)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -605,14 +281,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import psutil\n",
     "\n",
     "# You may change the settings in this cell according to Performance Test Tool result.\n",
-    "use_openmp = False\n",
+    "use_openmp = True\n",
     "\n",
     "# ATTENTION: these environment variables must be set before importing onnxruntime.\n",
     "if use_openmp:\n",
@@ -625,32 +301,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "PyTorch Inference time = 37.81 ms\n",
-      "OnnxRuntime Inference time = 31.81 ms\n",
-      "PyTorch and OnnxRuntime output 0 (last_state) are close: True\n",
-      "PyTorch and OnnxRuntime layer 0 state (present_0) are close: True\n",
-      "PyTorch and OnnxRuntime layer 1 state (present_1) are close: True\n",
-      "PyTorch and OnnxRuntime layer 2 state (present_2) are close: True\n",
-      "PyTorch and OnnxRuntime layer 3 state (present_3) are close: True\n",
-      "PyTorch and OnnxRuntime layer 4 state (present_4) are close: True\n",
-      "PyTorch and OnnxRuntime layer 5 state (present_5) are close: True\n",
-      "PyTorch and OnnxRuntime layer 6 state (present_6) are close: True\n",
-      "PyTorch and OnnxRuntime layer 7 state (present_7) are close: True\n",
-      "PyTorch and OnnxRuntime layer 8 state (present_8) are close: True\n",
-      "PyTorch and OnnxRuntime layer 9 state (present_9) are close: True\n",
-      "PyTorch and OnnxRuntime layer 10 state (present_10) are close: True\n",
-      "PyTorch and OnnxRuntime layer 11 state (present_11) are close: True\n",
-      "Wall time: 6.97 s\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import onnxruntime\n",
     "import numpy\n",
@@ -671,7 +324,7 @@
     "    sess_options.intra_op_num_threads=psutil.cpu_count(logical=True)\n",
     "\n",
     "# Specify providers when you use onnxruntime-gpu for CPU inference.\n",
-    "session = onnxruntime.InferenceSession(export_model_path, sess_options, providers=['CPUExecutionProvider'])\n",
+    "session = onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=['CPUExecutionProvider'])\n",
     "\n",
     "# Compare PyTorch and OnnxRuntime inference performance and results\n",
     "%time inference(model, session, input_ids, past=dummy_past if enable_past_input else None)"
@@ -679,20 +332,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "335"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
    "source": [
     "import gc\n",
     "del session\n",
@@ -701,7 +343,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -710,7 +352,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -719,47 +361,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "    BertOnnxModel.py: Fused LayerNormalization count: 25\n",
-      "    BertOnnxModel.py: Fused Gelu (FastGelu fits better) count: 12\n",
-      "    BertOnnxModel.py: Fused Reshape count:48\n",
-      "    BertOnnxModel.py: Fused SkipLayerNormalization count: 1\n",
-      "        OnnxModel.py: Removed unused constant nodes: 422\n",
-      "    BertOnnxModel.py: Fused Attention count:0\n",
-      "    BertOnnxModel.py: skip embed layer fusion since mask input is not found\n",
-      "    BertOnnxModel.py: opset verion: 11\n",
-      "        OnnxModel.py: Output model to ./gpt2_onnx\\gpt2_past0_optimized.onnx\n",
-      "    BertOnnxModel.py: EmbedLayer=0, Attention=0, Gelu=12, LayerNormalization=25, Successful=False\n",
-      "bert_model_optimization.py: The output model is not fully optimized. It might not be usable.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "# Local directory corresponding to https://github.com/microsoft/onnxruntime/tree/master/onnxruntime/python/tools/bert/\n",
-    "%run $bert_opt_script --input $export_model_path --output $optimized_model --opt_level 0"
+    "%run $bert_opt_script --model_type gpt2 --input $onnx_model_path --output $optimized_model --opt_level 0"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "PyTorch Inference time = 38.72 ms\n",
-      "OnnxRuntime Inference time = 29.58 ms\n",
-      "Wall time: 6.83 s\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "session = onnxruntime.InferenceSession(optimized_model, sess_options, providers=['CPUExecutionProvider'])\n",
     "\n",
@@ -782,56 +396,9 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": null,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{\n",
-      "  \"gpu\": {\n",
-      "    \"driver_version\": \"441.22\",\n",
-      "    \"devices\": [\n",
-      "      {\n",
-      "        \"memory_total\": 8589934592,\n",
-      "        \"memory_available\": 6569947136,\n",
-      "        \"name\": \"GeForce GTX 1070\"\n",
-      "      }\n",
-      "    ]\n",
-      "  },\n",
-      "  \"cpu\": {\n",
-      "    \"brand\": \"Intel(R) Core(TM) i7-8700 CPU @ 3.20GHz\",\n",
-      "    \"cores\": 6,\n",
-      "    \"logical_cores\": 12,\n",
-      "    \"hz\": \"3.1920 GHz\",\n",
-      "    \"l2_cache\": \"1536 KB\",\n",
-      "    \"l3_cache\": \"12288 KB\",\n",
-      "    \"processor\": \"Intel64 Family 6 Model 158 Stepping 10, GenuineIntel\"\n",
-      "  },\n",
-      "  \"memory\": {\n",
-      "    \"total\": 16971259904,\n",
-      "    \"available\": 2854060032\n",
-      "  },\n",
-      "  \"python\": \"3.6.10.final.0 (64 bit)\",\n",
-      "  \"os\": \"Windows-10-10.0.18362-SP0\",\n",
-      "  \"onnxruntime\": {\n",
-      "    \"version\": \"1.2.0\",\n",
-      "    \"support_gpu\": false\n",
-      "  },\n",
-      "  \"pytorch\": {\n",
-      "    \"version\": \"1.4.0+cpu\",\n",
-      "    \"support_gpu\": false\n",
-      "  },\n",
-      "  \"tensorflow\": {\n",
-      "    \"version\": \"2.1.0\",\n",
-      "    \"git_version\": \"v2.1.0-rc2-17-ge5bf8de410\",\n",
-      "    \"support_gpu\": true\n",
-      "  }\n",
-      "}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "machine_info_script = os.path.join(bert_tools_dir, 'MachineInfo.py')\n",
     "%run $machine_info_script --silent"

onnxruntime/python/tools/bert/notebooks/PyTorch_Bert-Squad_OnnxRuntime_GPU.ipynb

@@ -603,25 +603,25 @@
      ]
     }
    ],
-    "source": [
-      "import os\n",
-      "import wget\n",
-      "\n",
-      "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
-      "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
-      "\n",
-      "script_dir = './bert_scripts'\n",
-      "if not os.path.exists(script_dir):\n",
-      "    os.makedirs(script_dir)\n",
-      "\n",
-      "for filename in script_files:\n",
-      "    target_file = os.path.join(script_dir, filename)\n",
-      "    if enable_overwrite and os.path.exists(target_file):\n",
-      "        os.remove(target_file)\n",
-      "    if not os.path.exists(target_file):\n",
-      "        wget.download(url_prfix + filename, target_file)\n",
-      "        print(\"Downloaded\", filename)"
-    ]
+   "source": [
+    "import os\n",
+    "import wget\n",
+    "\n",
+    "url_prfix = \"https://raw.githubusercontent.com/microsoft/onnxruntime/master/onnxruntime/python/tools/bert/\"\n",
+    "script_files = ['bert_perf_test.py', 'bert_test_data.py', 'compare_bert_results.py', 'BertOnnxModel.py', 'BertOnnxModelKeras.py', 'BertOnnxModelTF.py', 'Gpt2OnnxModel.py', 'OnnxModel.py', 'bert_model_optimization.py', 'MachineInfo.py']\n",
+    "\n",
+    "script_dir = './bert_scripts'\n",
+    "if not os.path.exists(script_dir):\n",
+    "    os.makedirs(script_dir)\n",
+    "\n",
+    "for filename in script_files:\n",
+    "    target_file = os.path.join(script_dir, filename)\n",
+    "    if enable_overwrite and os.path.exists(target_file):\n",
+    "        os.remove(target_file)\n",
+    "    if not os.path.exists(target_file):\n",
+    "        wget.download(url_prfix + filename, target_file)\n",
+    "        print(\"Downloaded\", filename)"
+   ]
   },
   {
    "cell_type": "markdown",
@@ -676,9 +676,8 @@
     }
    ],
    "source": [
-    "GPU_OPTION = '--gpu_only' if use_gpu else ''\n",
     "optimized_fp32_model_path = './onnx/bert-base-cased-squad_opt_{}_fp32.onnx'.format('gpu' if use_gpu else 'cpu')\n",
-    "%run ./bert_scripts/bert_model_optimization.py --input $export_model_path --output $optimized_fp32_model_path $GPU_OPTION --input_int32"
+    "%run ./bert_scripts/bert_model_optimization.py --input $export_model_path --output $optimized_fp32_model_path --input_int32"
    ]
   },
   {
@@ -1148,9 +1147,8 @@
     }
    ],
    "source": [
-    "GPU_OPTION = '--gpu_only' if use_gpu else ''\n",
     "optimized_fp16_model_path = './onnx/bert-base-cased-squad_opt_{}_fp16.onnx'.format('gpu' if use_gpu else 'cpu')\n",
-    "%run  ./bert_scripts/bert_model_optimization.py --input $export_model_path --output $optimized_fp16_model_path $GPU_OPTION --float16 --input_int32"
+    "%run  ./bert_scripts/bert_model_optimization.py --input $export_model_path --output $optimized_fp16_model_path --float16 --input_int32"
    ]
   },
   {

onnxruntime/python/tools/bert/test_bert_optimization.py

@@ -129,16 +129,12 @@ class TestBertOptimization(unittest.TestCase):
         }
         self.verify_node_count(bert_model, expected_node_count)
 
-    def test_pytorch_model_0_cpu(self):
+    def test_pytorch_model_0(self):
         input = BERT_TEST_MODELS['bert_pytorch_0']
         bert_model = optimize_model(input,
                                     'bert',
-                                    gpu_only=False,
                                     num_heads=2,
-                                    hidden_size=8,
-                                    sequence_length=10,
-                                    input_int32=False,
-                                    float16=False)
+                                    hidden_size=8)
 
         expected_node_count = {
             'EmbedLayerNormalization': 1,
@@ -150,54 +146,21 @@ class TestBertOptimization(unittest.TestCase):
         }
         self.verify_node_count(bert_model, expected_node_count)
 
-    def test_pytorch_model_0_gpu(self):
-        if 'CUDAExecutionProvider' not in onnxruntime.get_available_providers():
-            print("skip test_pytorch_model_0_gpu since no gpu found")
-            return
-
-        input = BERT_TEST_MODELS['bert_pytorch_0']
-        bert_model = optimize_model(input,
-                                    'bert',
-                                    gpu_only=True,
-                                    num_heads=2,
-                                    hidden_size=8,
-                                    sequence_length=10,
-                                    input_int32=False,
-                                    float16=False)
-
-        expected_node_count = {
-            'EmbedLayerNormalization': 1,
-            'Attention': 12,
-            'SkipLayerNormalization': 24,
-            'FastGelu': 12,
-            'Gelu': 0,
-            'BiasGelu': 0
-        }
-        self.verify_node_count(bert_model, expected_node_count)
-
-    def test_pytorch_model_2_cpu(self):
+    def test_pytorch_model_2(self):
         input = BERT_TEST_MODELS['bert_squad_pytorch1.4_opset10_fp32']
         bert_model = optimize_model(input,
                                     'bert',
-                                    gpu_only=False,
                                     num_heads=2,
-                                    hidden_size=8,
-                                    sequence_length=10,
-                                    input_int32=False,
-                                    float16=False)
+                                    hidden_size=8)
         self.assertTrue(bert_model.is_fully_optimized())
 
-    def test_keras_model_1_cpu(self):
+    def test_keras_model_1(self):
         input = BERT_TEST_MODELS['bert_keras_0']
 
         bert_model = optimize_model(input,
                                     'bert_keras',
-                                    gpu_only=False,
                                     num_heads=2,
-                                    hidden_size=8,
-                                    sequence_length=7,
-                                    input_int32=False,
-                                    float16=False)
+                                    hidden_size=8)
 
         expected_node_count = {
             'EmbedLayerNormalization': 1,
@@ -210,17 +173,13 @@ class TestBertOptimization(unittest.TestCase):
         }
         self.verify_node_count(bert_model, expected_node_count)
 
-    def test_keras_squad_model_cpu(self):
+    def test_keras_squad_model(self):
         input = BERT_TEST_MODELS['bert_keras_squad']
 
         bert_model = optimize_model(input,
                                     'bert_keras',
-                                    gpu_only=False,
                                     num_heads=2,
-                                    hidden_size=8,
-                                    sequence_length=7,
-                                    input_int32=False,
-                                    float16=False)
+                                    hidden_size=8)
 
         self.assertTrue(bert_model.is_fully_optimized())
 
@@ -228,12 +187,8 @@ class TestBertOptimization(unittest.TestCase):
         input = BERT_TEST_MODELS['gpt2']
         bert_model = optimize_model(input,
                                     'gpt2',
-                                    gpu_only=False,
                                     num_heads=2,
-                                    hidden_size=4,
-                                    sequence_length=2,
-                                    input_int32=False,
-                                    float16=False)
+                                    hidden_size=4)
 
         expected_node_count = {
             'EmbedLayerNormalization': 0,

onnxruntime/test/framework/allocation_planner_test.cc

@@ -166,7 +166,7 @@ class PlannerTest : public ::testing::Test {
   PlannerTest()
       : model_("test", false, DefaultLoggingManager().DefaultLogger()),
         graph_(model_.MainGraph()),
-        tp_(concurrency::CreateThreadPool(&onnxruntime::Env::Default(), OrtThreadPoolParams())),
+        tp_(concurrency::CreateThreadPool(&onnxruntime::Env::Default(), OrtThreadPoolParams(), concurrency::ThreadPoolType::INTRA_OP)),
         state_(execution_providers_, false, tp_.get(), nullptr) {
     std_kernel_ = KernelDefBuilder().SetName("Transpose").Provider(kCpuExecutionProvider).SinceVersion(1, 10).Build();
     in_place_kernel_ =
@@ -201,8 +201,8 @@ class PlannerTest : public ::testing::Test {
 
   void BindKernel(onnxruntime::Node* p_node, ::onnxruntime::KernelDef& kernel_def, KernelRegistry* reg) {
     auto info = onnxruntime::make_unique<OpKernelInfo>(*p_node, kernel_def, *execution_providers_.Get(*p_node),
-                                               state_.GetInitializedTensors(), state_.GetOrtValueNameIdxMap(),
-                                               state_.GetFuncMgr(), state_.GetDataTransferMgr());
+                                                       state_.GetInitializedTensors(), state_.GetOrtValueNameIdxMap(),
+                                                       state_.GetFuncMgr(), state_.GetDataTransferMgr());
     op_kernel_infos_.push_back(std::move(info));
     if (reg->TryFindKernel(*p_node, onnxruntime::kCpuExecutionProvider) == nullptr) {
       auto st = reg->Register(

onnxruntime/test/framework/math_test.cc

@@ -28,12 +28,12 @@ namespace onnxruntime {
 //parameter is thread pool size
 class MathGemmTest : public testing::TestWithParam<int> {
  protected:
-   static OrtThreadPoolParams CreateThreadPoolOptions(int size){
-       OrtThreadPoolParams option;
-       option.thread_pool_size = size;
-       return option;
-   }
-  std::unique_ptr<concurrency::ThreadPool> tp{concurrency::CreateThreadPool(&Env::Default(),CreateThreadPoolOptions(GetParam()))};
+  static OrtThreadPoolParams CreateThreadPoolOptions(int size) {
+    OrtThreadPoolParams option;
+    option.thread_pool_size = size;
+    return option;
+  }
+  std::unique_ptr<concurrency::ThreadPool> tp{concurrency::CreateThreadPool(&Env::Default(), CreateThreadPoolOptions(GetParam()), concurrency::ThreadPoolType::INTRA_OP)};
 };
 
 TEST_P(MathGemmTest, GemmNoTransNoTrans) {
@@ -124,7 +124,7 @@ TEST_P(MathGemmTest, GemmNoTransTrans) {
 }
 
 INSTANTIATE_TEST_SUITE_P(MathGemmTests, MathGemmTest,
-                        testing::Values(1, 0));
+                         testing::Values(1, 0));
 
 TEST(MathTest, GemvNoTrans) {
   auto& provider = CPUMathUtil::Instance();

onnxruntime/test/framework/session_state_test.cc

@@ -42,7 +42,7 @@ class SessionStateAddGetKernelTest : public testing::TestWithParam<int> {};
 TEST_P(SessionStateAddGetKernelTest, AddGetKernelTest) {
   OrtThreadPoolParams to;
   to.thread_pool_size = GetParam();
-  auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to);
+  auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP);
   ONNX_OPERATOR_SCHEMA(Variable)
       .SetDoc("Input variable.")
       .Output(0, "output_1", "docstr for output_1.", "tensor(int32)");
@@ -96,8 +96,7 @@ class TestParam {
   bool enable_mem_pattern;
   int thread_count;
 };
-TestParam param_list[] = {{3, true, 0}, {4, true, 0}, {3, false, 0}, {4, false, 0},
-                          {3, true, 1}, {4, true, 1}, {3, false, 1}, {4, false, 1}};
+TestParam param_list[] = {{3, true, 0}, {4, true, 0}, {3, false, 0}, {4, false, 0}, {3, true, 1}, {4, true, 1}, {3, false, 1}, {4, false, 1}};
 }  // namespace
 class SessionStateTestP : public testing::TestWithParam<TestParam> {};
 // Test that we separate out constant and non-constant initializers correctly
@@ -105,7 +104,7 @@ TEST_P(SessionStateTestP, TestInitializerProcessing) {
   const TestParam& param = GetParam();
   OrtThreadPoolParams to;
   to.thread_pool_size = to.thread_pool_size;
-  auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to);
+  auto tp = concurrency::CreateThreadPool(&onnxruntime::Env::Default(), to, concurrency::ThreadPoolType::INTRA_OP);
 
   std::basic_ostringstream<ORTCHAR_T> oss;
   oss << ORT_TSTR("testdata/optional_inputs_ir") << param.ir_version << ORT_TSTR(".onnx");

onnxruntime/test/providers/cpu/math/topk_op_test.cc

@@ -425,54 +425,17 @@ TEST(TopKOperator, Top1ExplicitAxisMultiDInputSmallestElements) {
   top_1_explicit_axis_MultiD_input_smallest(11, 0);  //unsorted
 }
 
-TEST(TopKOperator, SelectFirstSortNext) {
-  // in this test, we will select the top 5 elements first then sort the chosen 5 elements
-  // Select + Sort  = O(n + k * ln(k)) = 50 + 5 * ln(5) = 58.047
-  // Sorted selection: O(n * ln(k)) = 50 * ln(5) = 80.47
-  // The algorithm used will be Select + Sort
-  std::vector<float> input_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0,
-                                   11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 16.0f, 17.0f, 18.0f, 19.0f, 20.0,
-                                   21.0f, 22.0f, 23.0f, 24.0f, 25.0f, 26.0f, 27.0f, 28.0f, 29.0f, 30.0,
-                                   31.0f, 32.0f, 33.0f, 34.0f, 35.0f, 36.0f, 37.0f, 38.0f, 39.0f, 40.0,
-                                   41.0f, 42.0f, 43.0f, 44.0f, 45.0f, 46.0f, 47.0f, 48.0f, 49.0f, 50.0};
-  std::vector<int64_t> input_dimensions = {50};
-  std::vector<float> expected_vals = {50.0f, 49.0f, 48.0f, 47.0f, 46.0f};
-  std::vector<int64_t> expected_indices = {49, 48, 47, 46, 45};
-  std::vector<int64_t> expected_dimensions = {5};
-  int64_t axis = 0;
-  RunTest(11, 5, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis);  // largest values
-}
-
-TEST(TopKOperator, SelectFirstSortNextInt64) {
-  // in this test, we will select the top 5 elements first then sort the chosen 5 elements
-  // Select + Sort  = O(n + k * ln(k)) = 50 + 5 * ln(5) = 58.047
-  // Sorted selection: O(n * ln(k)) = 50 * ln(5) = 80.47
-  // The algorithm used will be Select + Sort
-  std::vector<int64_t> input_vals = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
-                                     11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
-                                     21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
-                                     31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
-                                     41, 42, 43, 44, 45, 46, 47, 48, 49, 50};
-  std::vector<int64_t> input_dimensions = {50};
-  std::vector<int64_t> expected_vals = {50, 49, 48, 47, 46};
-  std::vector<int64_t> expected_indices = {49, 48, 47, 46, 45};
-  std::vector<int64_t> expected_dimensions = {5};
-  int64_t axis = 0;
-  RunTest(11, 5, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis);  // largest values
-}
-
-TEST(TopKOperator, SortedSelection) {
-  // in this test, we will use sorted selection (using heap)
-  // Select + Sort  = O(n + k * ln(k)) = 10 + 5 * ln(5) = 18.04
-  // Sorted selection: O(n * ln(k)) = 10 * ln(5) = 16.09
-  // The algorithm used will be Sorted selection
-  std::vector<float> input_vals = {10.0f, 8.0f, 7.0f, 4.0f, 5.0f, 6.0f, 1.0f, 2.0f, 9.0f, 3.0};
-  std::vector<int64_t> input_dimensions = {10};
-  std::vector<float> expected_vals = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
-  std::vector<int64_t> expected_indices = {6, 7, 9, 3, 4};
-  std::vector<int64_t> expected_dimensions = {5};
-  int64_t axis = 0;
-  RunTest(11, 5, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, axis, 0);  // smallest values
+// test path where SelectTopK is used (select using std::nth_element)
+// we use a custom path for n=1, and priority queue based implementation if
+//   bool use_priority_queue = k != 1 && (k < 4 || (std::log2(k) / std::log2(n)) < 0.725);
+// so easiest way to test is for k to be 4 and n to be a little larger
+TEST(TopKOperator, NthElement) {
+  std::vector<float> input_vals = {10.0f, 8.0f, 7.0f, 4.0f, 5.0f, 6.0f};
+  std::vector<int64_t> input_dimensions = {6};
+  std::vector<float> expected_vals = {10.0f, 8.0f, 7.0f, 6.0f};
+  std::vector<int64_t> expected_indices = {0, 1, 2, 5};
+  std::vector<int64_t> expected_dimensions = {4};
+  RunTest(11, 4, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
 }
 
 // test dimension in range (GridDim::maxThreadsPerBlock, GridDim::maxThreadsPerBlock * 2], ie. [257, 512]
@@ -532,5 +495,66 @@ TEST(TopKOperator, BigArrayBigTopKSorted) {
   RunTest(11, 9000, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, 0, 1, 1);
 }
 
+static void top_3_all_same(int opset_version, int64_t largest = 1) {
+  // whether it's largest or smallest we should pick the first instance/s of a number if there are multiple
+  std::vector<float> input_vals = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f};
+  std::vector<int64_t> input_dimensions = {2, 4};
+  std::vector<float> expected_vals = {0.1f, 0.1f, 0.1f, 0.1f, 0.1f, 0.1f};
+  std::vector<int64_t> expected_indices = {0, 1, 2, 0, 1, 2};
+  std::vector<int64_t> expected_dimensions = {2, 3};
+  RunTest(opset_version, 3, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false, -1, largest);
+}
+
+TEST(TopKOperator, Top3AllSame) {
+  top_3_all_same(10);
+  top_3_all_same(11);
+  top_3_all_same(10, 0);  // smallest
+  top_3_explicit_axis(11, 0);
+}
+
+static void TestThreaded(int64_t k, int64_t n, int64_t batch_size) {
+  std::vector<float> input_vals(n * batch_size, 0.0f);
+  std::iota(input_vals.begin(), input_vals.end(), 0.0f);
+
+  std::vector<int64_t> input_dimensions = {n, batch_size};
+
+  std::vector<float> expected_vals(n * k, 0.0f);
+  std::vector<int64_t> expected_indices(n * k, 0);
+  std::vector<int64_t> expected_dimensions = {n, k};
+
+  for (int64_t i = 0; i < n; ++i) {
+    auto begin_batch_output = expected_vals.begin() + i * k;
+    std::iota(begin_batch_output, begin_batch_output + k, static_cast<float>(((i + 1) * batch_size) - k));
+    std::reverse(begin_batch_output, begin_batch_output + k);
+
+    // indices are within the axis so don't need adjusting by the batch number
+    auto begin_indices_output = expected_indices.begin() + i * k;
+    std::iota(begin_indices_output, begin_indices_output + k, batch_size - k);
+    std::reverse(begin_indices_output, begin_indices_output + k);
+  }
+
+  RunTest(11, k, input_vals, input_dimensions, expected_vals, expected_indices, expected_dimensions, false);
+}
+
+// create input of 2x1000 and select 200 so 2 threads are needed based on there being 2 rows
+// and sufficient items to process given this calculation:
+//   int64_t threads_needed = static_cast<int64_t>(std::floor(input_shape.Size() * k / (128 * 1024)));
+TEST(TopKOperator, PriorityQueueThreaded) {
+  const int64_t k = 200;
+  const int64_t n = 2;
+  const int64_t batch_size = 1000;
+  TestThreaded(k, n, batch_size);
+}
+
+// create input of 2x500 and select 400 so 2 threads are needed based on there being 2 rows
+// and sufficient items to process given this calculation:
+//   int64_t threads_needed = static_cast<int64_t>(std::floor(input_shape.Size() * k / (128 * 1024)));
+TEST(TopKOperator, SelectTopKThreaded) {
+  const int64_t k = 400;
+  const int64_t n = 2;
+  const int64_t batch_size = 500;
+  TestThreaded(k, n, batch_size);
+}
+
 }  // namespace test
 }  // namespace onnxruntime