Added xsmm backend for torch_mlir compiler (#70)

Signed-off-by: Gregory Shimansky <[email protected]>

Author: gshimansky

.github/workflows/test-single-config.yml

@@ -21,6 +21,7 @@ on:
                     - torch
                     - dynamo
                     - torch_mlir
+                    - torch_mlir_xsmm
                     - torchscript
                     - torchscript_onednn
                     - ipex

.github/workflows/test.yml

@@ -40,7 +40,8 @@ jobs:
                   {device: 'cpu', compiler: 'ipex'},
                   {device: 'cpu', compiler: 'ipex_onednn_graph'},
 #                  {device: 'xpu', compiler: 'ipex'},
-                  {device: 'cpu', compiler: 'torch_mlir'}
+                  {device: 'cpu', compiler: 'torch_mlir'},
+                  {device: 'cpu', compiler: 'torch_mlir_xsmm'}
                 ]
                 test_script: ${{ fromJson(inputs.test_scripts) }}
             fail-fast: false
@@ -56,12 +57,3 @@ jobs:
             test_script: ${{ matrix.test_script }}
         secrets:
             DB_URL: ${{ secrets.DB_URL }}
-
-    shutdown:
-        needs: mlp_test
-        if: ${{ contains(inputs.runner_type, 'amd') }} && inputs.shutdown_cloud_runner
-        runs-on: ${{ inputs.runner_type }}
-        steps:
-            - name: shutdown
-              shell: bash -el {0}
-              run: sudo shutdown -h +2

dl_bench/cli/launcher.py

@@ -85,6 +85,7 @@ def parse_args():
             "ipex",
             "ipex_onednn_graph",
             "torch_mlir",
+            "torch_mlir_xsmm",
         ],
         help="Compilation mode to use. No compilation by default.",
     )

dl_bench/utils.py

@@ -268,7 +268,7 @@ def _compile_model(compile_mode: str, device, model: Module, sample_input, dtype
 
             compiled_model = dynamo.optimize(be.refbackend_torchdynamo_backend)(model)
             print("Compiled with torch_mlir (torchscript, inference)")
-        elif compile_mode == "torch_mlir":
+        elif compile_mode == "torch_mlir" or compile_mode == "torch_mlir_xsmm":
             from torch_mlir._dynamo_fx_importer import import_fx_graph_as_func
             from torch_mlir_e2e_test.configs.torchdynamo import jit
             from torch_mlir_e2e_test.framework import TestOptions
@@ -277,6 +277,9 @@ def _compile_model(compile_mode: str, device, model: Module, sample_input, dtype
             from torch_mlir_e2e_test.linalg_on_tensors_backends.cpuprotobackend import (
                 CpuProtoLinalgOnTensorsBackend,
             )
+            from torch_mlir_e2e_test.linalg_on_tensors_backends.xsmmprotobackend import (
+                XsmmProtoLinalgOnTensorsBackend,
+            )
             import torch.utils._pytree as pytree
 
             # debug_timer seems to cause problems:
@@ -290,7 +293,7 @@ def _compile_model(compile_mode: str, device, model: Module, sample_input, dtype
                 opts,
                 output_type="linalg-on-tensors",
             )
-            backend = CpuProtoLinalgOnTensorsBackend(opts)
+            backend = CpuProtoLinalgOnTensorsBackend(opts) if compile_mode == "torch_mlir" else XsmmProtoLinalgOnTensorsBackend(opts)
             # backend = RefBackendLinalgOnTensorsBackend()
             module = backend.compile(module)
             backend_module = backend.load(module)