Skip to content

Instantly share code, notes, and snippets.

from shark.shark_inference import SharkInference
from shark.shark_downloader import download_torch_model
import numpy as np
mlir_model, func_name, inputs, golden_out = download_torch_model("v_diffusion")
shark_module = SharkInference(
mlir_model, func_name, device="cpu", mlir_dialect="linalg"
)
This file has been truncated, but you can view the full file.
#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d0, 0)>
#map2 = affine_map<(d0, d1) -> ()>
#map3 = affine_map<(d0, d1) -> (d1, d0)>
#map4 = affine_map<(d0, d1) -> (d1)>
#map5 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (0, d1, 0, d3, 0, 0, 0, 0)>
#map6 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>
#map7 = affine_map<(d0, d1, d2, d3) -> (d1)>
#map8 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map9 = affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, 0)>

Python Script

Script can be found at examples/onnx_add.py

import torch
import torch_mlir
from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend

class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()

Python Script

Script can be found at examples/onnx_resnet18.py

import sys 

from PIL import Image
import requests

import torch
import torchvision.models as models

Python Script

Script can be found at examples/onnx_logsoftmax.py

import torch
import torch_mlir

class ToyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1) -> (d1, d0)>
#map2 = affine_map<(d0, d1) -> (0, d1)>
#map3 = affine_map<(d0, d1) -> (d1)>
module attributes {torch.debug_module_name = "DLRM_Net"} {
func.func @forward(%arg0: tensor<1x4xf32>, %arg1: tensor<3x1xi64>, %arg2: tensor<3xi64>, %arg3: tensor<1xi64>, %arg4: tensor<1xi64>) -> tensor<1x2xf32> {
%cst = arith.constant dense<[[0.929304063, 0.0979973599, 0.239170983], [-5.614850e-01, -1.25276566, -0.220038965]]> : tensor<2x3xf32>
%cst_0 = arith.constant dense<[0.110555418, 0.869946897]> : tensor<2xf32>
%cst_1 = arith.constant dense<[[0.237254873, 0.178356424, 0.798618853, -0.109661706], [0.167341724, -0.456533372, -1.36463046, 0.349373847], [0.462060571, -0.396703899, 1.2132349, -0.777391135]]> : tensor<3x4xf32>
%cst_2 = arith.constant dense<[0.0264186915, -0.108070649, 0.884950518]> : tensor<3xf32>
@qedawkins
qedawkins / copy.cpp
Created September 7, 2022 18:14
Single Thread Level-Zero Dispatch
// Example for dispatching a SPIR-V Kernel using Level Zero on the Intel HD Graphics
// Sample based on the test-suite exanples from Level-Zero:
// https://github.com/intel/compute-runtime/blob/master/level_zero/core/test/black_box_tests/zello_world_gpu.cpp
#include "ze_api.h"
#include <chrono>
#include <cstring>
#include <fstream>
#include <iostream>

Update

IR for the embedding bag only is below

#map0 = affine_map<(d0, d1, d2) -> (d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d2)>
module attributes {torch.debug_module_name = "ToyEmbeddingBag"} {
  func.func @forward(%arg0: tensor<4xi64>, %arg1: tensor<2xi64>) -> tensor<2x3xf32> {
    %cst = arith.constant dense<[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0], [13.0, 14.0, 15.0], [16.0, 17.0, 18.0], [19.0, 20.0, 21.0], [22.0, 23.0, 24.0], [25.0, 26.0, 27.0], [28.0, 29.0, 30.0]]> : tensor<10x3xf32>
    %cst_0 = arith.constant 0.000000e+00 : f32