Merge branch 'develop' into migraphx-support-for-gfx12

Jenkinsfile

@@ -1,3 +1,7 @@
+import org.jenkinsci.plugins.pipeline.modeldefinition.Utils
+
+DOCKER_IMAGE = 'rocm/migraphx-ci-jenkins-ubuntu'
+
 def getgputargets() {
     targets="gfx906;gfx908;gfx90a;gfx1030;gfx1100;gfx1101;gfx1102"
     return targets
@@ -60,19 +64,16 @@ def rocmtestnode(Map conf) {
                 sh 'printenv'
                 checkout scm
             }
-            gitStatusWrapper(credentialsId: "${env.migraphx_ci_creds}", gitHubContext: "Jenkins - ${variant}", account: 'ROCmSoftwarePlatform', repo: 'AMDMIGraphX') {
-                pre()
-                stage("image ${variant}") {
-                    try {
-                        docker.build("${image}", "${docker_build_args} .")
-                    } catch(Exception ex) {
-                        docker.build("${image}", "${docker_build_args} --no-cache .")
 
-                    }
-                }
-                withDockerContainer(image: image, args: "--device=/dev/kfd --device=/dev/dri --group-add video --cap-add SYS_PTRACE -v=/home/jenkins:/home/jenkins ${docker_args}") {
-                    timeout(time: 2, unit: 'HOURS') {
-                        body(cmake_build)
+            gitStatusWrapper(credentialsId: "${env.migraphx_ci_creds}", gitHubContext: "Jenkins - ${variant}", account: 'ROCmSoftwarePlatform', repo: 'AMDMIGraphX') {
+                withCredentials([usernamePassword(credentialsId: 'docker_test_cred', passwordVariable: 'DOCKERHUB_PASS', usernameVariable: 'DOCKERHUB_USER')]) {
+                    sh "echo $DOCKERHUB_PASS | docker login --username $DOCKERHUB_USER --password-stdin"
+                    pre()
+                    sh "docker pull ${DOCKER_IMAGE}:${env.IMAGE_TAG}"
+                    withDockerContainer(image: "${DOCKER_IMAGE}:${env.IMAGE_TAG}", args: "--device=/dev/kfd --device=/dev/dri --group-add video --cap-add SYS_PTRACE -v=/home/jenkins/:/home/jenkins ${docker_args}") {
+                        timeout(time: 2, unit: 'HOURS') {
+                            body(cmake_build)
+                        }
                     }
                 }
             }
@@ -118,6 +119,47 @@ def rocmnode(name, body) {
     }
 }
 
+properties([
+    parameters([
+        booleanParam(name: 'FORCE_DOCKER_IMAGE_BUILD', defaultValue: false)
+    ])
+])
+
+node("(rocmtest || migraphx)") {
+    Boolean imageExists = false
+    withCredentials([usernamePassword(credentialsId: 'docker_test_cred', passwordVariable: 'DOCKERHUB_PASS', usernameVariable: 'DOCKERHUB_USER')]) {
+        sh "echo $DOCKERHUB_PASS | docker login --username $DOCKERHUB_USER --password-stdin"
+        stage('Check image') {
+            sh 'printenv'
+            checkout scm
+            def calculateImageTagScript = """
+                shopt -s globstar
+                sha256sum **/Dockerfile **/*requirements.txt **/install_prereqs.sh **/rbuild.ini | sha256sum | cut -d " " -f 1
+            """
+            env.IMAGE_TAG = sh(script: "bash -c '${calculateImageTagScript}'", returnStdout: true).trim()
+            imageExists = sh(script: "docker manifest inspect ${DOCKER_IMAGE}:${IMAGE_TAG}", returnStatus: true) == 0
+        }
+        stage('Build image') {
+            if(!imageExists || params.FORCE_DOCKER_IMAGE_BUILD) {
+                def builtImage
+
+                try {
+                    sh "docker pull ${DOCKER_IMAGE}:latest"
+                    builtImage = docker.build("${DOCKER_IMAGE}:${IMAGE_TAG}", "--cache-from ${DOCKER_IMAGE}:latest .")
+                } catch(Exception ex) {
+                    builtImage = docker.build("${DOCKER_IMAGE}:${IMAGE_TAG}", " --no-cache .")
+                }
+                builtImage.push("${IMAGE_TAG}")
+                builtImage.push("latest")
+            } else {
+                echo "Image already exists, skip building available"
+                // Skip stage so it remains in the visualization
+                Utils.markStageSkippedForConditional(STAGE_NAME)
+            }
+        }
+    }
+}
+
 rocmtest clang_debug: rocmnode('mi100+') { cmake_build ->
     stage('hipRTC Debug') {
         // Disable MLIR since it doesnt work with all ub sanitizers

docs/dev/env_vars.rst

@@ -126,6 +126,10 @@ Use it in conjunction with ``MIGRAPHX_DISABLE_MLIR=1``.
 Set to "1", "enable", "enabled", "yes", or "true" to use.
 Disables use of the rocMLIR library.
 
+.. envvar:: MIGRAPHX_ENABLE_HIPBLASLT_GEMM
+Set to "1", "enable", "enabled", "yes", or "true" to use.
+Enables use of hipBLASLt.
+
 .. envvar:: MIGRAPHX_COPY_LITERALS
 
 Set to "1", "enable", "enabled", "yes", or "true" to use.
@@ -299,6 +303,11 @@ Enable reduction fusions in MLIR.
 Set to "1", "enable", "enabled", "yes", or "true" to use.
 Enable Split-k perf configs when tuning with MLIR.
 
+.. envvar:: MIGRAPHX_MLIR_DUMP_TO_MXR
+
+Set to path where MXRs will be saved.
+Dumps MLIRs module to mxr files.
+
 CK vars
 -----------
 
@@ -317,6 +326,14 @@ Mandatorily adds ``-DMIGRAPHX_CK_CHECK=1`` for compiling composable kernel opera
 Set to "1", "enable", "enabled", "yes", or "true" to use.
 Performs tuning for composable kernels.
 
+hipBLASLt vars
+--------------
+
+.. envvar:: MIGRAPHX_ENABLE_HIP_GEMM_TUNING
+
+Set to "1", "enable", "enabled", "yes", or "true" to use.
+Performs exhaustive tuning for hipBLASLt.
+
 Testing 
 ------------
 

docs/sphinx/requirements.in

@@ -1,2 +1,2 @@
-rocm-docs-core==1.8.2
+rocm-docs-core==1.8.3
 sphinx-collapse

docs/sphinx/requirements.txt

@@ -116,7 +116,7 @@ requests==2.32.3
     # via
     #   pygithub
     #   sphinx
-rocm-docs-core==1.8.2
+rocm-docs-core==1.8.3
     # via -r requirements.in
 smmap==5.0.1
     # via gitdb

examples/nlp/python_rnnt/README.md

@@ -18,7 +18,7 @@ python3 -m venv rnnt_venv
 3) Install dependencies
 
 ```bash
-pip install -r requirements.txt
+pip install -r torch_requirements.txt -r requirements.txt
 ```
 
 4) Download helper files 

examples/nlp/python_rnnt/requirements.txt

@@ -21,15 +21,9 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 #####################################################################################
--f https://repo.radeon.com/rocm/manylinux/rocm-rel-6.0/
-torch==2.1.1
-accelerate
-diffusers
-optimum[onnxruntime]
-transformers
-
 --extra-index-url https://test.pypi.org/simple
 hip-python
 
 librosa
-soundfile
+soundfile
+datasets

examples/nlp/python_rnnt/rnnt.py

@@ -23,10 +23,11 @@
 from typing import List, Optional, Tuple
 import argparse
 
+import torch
 import migraphx as mgx
 import os
 import sys
-import torch
+
 import torch.nn.functional as F
 
 from rnnt_data import librespeech_huggingface
@@ -88,9 +89,12 @@ def allocate_torch_tensors(model):
 
 
 class RNNT_MGX():
-    def __init__(self, seq_length, onnx_model_path):
+    def __init__(self, x, onnx_model_path, fp16):
 
-        fp16 = []
+        if fp16 is None:
+            fp16 = []
+        elif "all" in fp16:
+            fp16 = ["rnnt_encoder", "rnnt_prediction", "rnnt_joint"]
 
         compiled_model_path = None
         force_compile = False
@@ -102,8 +106,7 @@ def __init__(self, seq_length, onnx_model_path):
             RNNT_MGX.load_mgx_model(
                 "rnnt_encoder",
                 {
-                    "input": [seq_length, 1, 240
-                              ],  # seq_length, batch_size, feature_length
+                    "input": x.shape,  # seq_length, batch_size, feature_length
                     "feature_length": [1]
                 },
                 onnx_model_path,
@@ -334,6 +337,13 @@ def decode_string(result):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('--onnx_model_path', default='models/rnnt/')
+    parser.add_argument(
+        "--fp16",
+        choices=["all", "rnnt_encoder", "rnnt_prediction", "rnnt_joint"],
+        nargs="+",
+        help="Quantize models with fp16 precision.",
+    )
+
     args = parser.parse_args()
 
     model_parts = ['rnnt_encoder', 'rnnt_joint', 'rnnt_prediction']
@@ -344,7 +354,6 @@ def decode_string(result):
 
     print("Getting data...")
     x, out_lens, transcript = librespeech_huggingface()
-    seq_length = x.shape[0]
 
     if any(check_onnx_files) == False:
         from rnnt_torch_model import pytorch_rnnt_model
@@ -354,12 +363,13 @@ def decode_string(result):
         pytorch_model = pytorch_rnnt_model()
 
         print("Export pytorch model to ONNX...")
-        export_rnnt_onnx(pytorch_model, seq_length=seq_length)
+        export_rnnt_onnx(pytorch_model, x, args.onnx_model_path)
 
     print("Read MIGX model from ONNX and run...")
-    migx_model = RNNT_MGX(seq_length=seq_length,
-                          onnx_model_path=args.onnx_model_path)
-    rnnt_migx = GreedyDecoder(migx_model)
-    _, _, result = rnnt_migx.run(x.to(torch.float32), out_lens)
+    migx_model = RNNT_MGX(x,
+                          onnx_model_path=args.onnx_model_path,
+                          fp16=args.fp16)
+    rnnt_decoder = GreedyDecoder(migx_model)
+    _, _, result = rnnt_decoder.run(x.to(torch.float32), out_lens)
     print("Transcribed Sentence: ", decode_string(result))
     print("Ground Truth: ", transcript)

examples/nlp/python_rnnt/rnnt_onnx.py

@@ -7,13 +7,12 @@ def make_directories(dir_path, model_name):
         os.makedirs(f"{dir_path}/{model_name}/")
 
 
-def export_rnnt_onnx(model, seq_length):
-    dir_path = './models/rnnt/'
+def export_rnnt_onnx(model, x, dir_path):
 
     for component in ['rnnt_encoder', 'rnnt_prediction', 'rnnt_joint']:
         make_directories(dir_path, component)
 
-    batch_size, feature_length = 1, 240
+    seq_length, batch_size, feature_length = x.shape
     inp = torch.randn([seq_length, batch_size, feature_length]).to('cuda')
     feature_length = torch.LongTensor([seq_length]).to('cuda')
     x_padded, x_lens = model.encoder(inp, feature_length)

examples/nlp/python_rnnt/torch_requirements.txt

@@ -0,0 +1,2 @@
+--index-url https://download.pytorch.org/whl/rocm6.1/
+torch

requirements.txt

@@ -28,4 +28,4 @@ pybind/pybind11@3e9dfa2866941655c56877882565e7577de6fc7b --build
 msgpack/[email protected] -DMSGPACK_BUILD_TESTS=Off
 sqlite3@3.43.2 -DCMAKE_POSITION_INDEPENDENT_CODE=On
 ROCm/composable_kernel@57cdd70b7cb14e5e3b60cd9a5f96ba8dc343763e -DCK_BUILD_JIT_LIB=On -DCMAKE_POSITION_INDEPENDENT_CODE=On
-ROCm/rocMLIR@cd281cf440ac52b94eff7fca5d98b04445ccc82b -DBUILD_FAT_LIBROCKCOMPILER=On
+ROCm/rocMLIR@e454b5d06fc2f099f7de3ee43450e7a6b1efe015 -DBUILD_FAT_LIBROCKCOMPILER=On

src/CMakeLists.txt

@@ -174,6 +174,7 @@ register_migraphx_ops(
     gathernd
     get_tuple_elem
     greater
+    group_query_attention
     gru
     identity
     if_op
@@ -383,5 +384,5 @@ rocm_export_targets(
     Threads
     ${MIGRAPHX_CONFIG_DEPENDS}
 )
- 
+