mlir/test/Dialect/MLProgram/ops.mlir

   1 // RUN: mlir-opt %s --allow-unregistered-dialect | mlir-opt --allow-unregistered-dialect | FileCheck %s
   2 // RUN: mlir-opt %s --allow-unregistered-dialect --mlir-print-op-generic | mlir-opt --allow-unregistered-dialect | FileCheck %s
   3
   4 // CHECK-LABEL: ml_program.func private @extern_func
   5 ml_program.func private @extern_func(i32) -> i32
   6
   7 // CHECK-LABEL: ml_program.func @defined_func
   8 ml_program.func @defined_func(%arg0 : i32) -> i32 {
   9   ml_program.return %arg0 : i32
  10 }
  11
  12 // CHECK-LABEL: ml_program.subgraph private @extern_subgraph
  13 ml_program.subgraph private @extern_subgraph(i32) -> i32
  14
  15 // CHECK-LABEL: ml_program.subgraph @compute_subgraph
  16 ml_program.subgraph @compute_subgraph(%arg0 : i32) -> i32 {
  17   %token = ml_program.token
  18   %1 = "unregistered.dummy"(%0, %token) : (i32, !ml_program.token) -> i32
  19   %0 = "unregistered.dummy"(%arg0) : (i32) -> i32
  20   ml_program.output %0 : i32
  21 }
  22
  23 // CHECK: ml_program.global private @global_same_type(dense<4> : tensor<4xi32>) : tensor<4xi32>
  24 ml_program.global private @global_same_type(dense<4> : tensor<4xi32>) : tensor<4xi32>
  25
  26 // CHECK: ml_program.global private mutable @global_mutable_undef : tensor<?xi32>
  27 ml_program.global private mutable @global_mutable_undef : tensor<?xi32>
  28
  29 // CHECK: ml_program.global private mutable @global_extern(#extern) : tensor<?xi32>
  30 ml_program.global private mutable @global_extern(#ml_program.extern : tensor<4xi32>) : tensor<?xi32>
  31
  32 // CHECK-LABEL: @global_load_const
  33 ml_program.func @global_load_const() -> tensor<4xi32> {
  34   %0 = ml_program.global_load_const @global_same_type : tensor<4xi32>
  35   ml_program.return %0 : tensor<4xi32>
  36 }
  37
  38 // CHECK-LABEL: @global_load_store
  39 ml_program.func @global_load_store() {
  40   %0 = ml_program.global_load @global_mutable_undef : tensor<?xi32>
  41   ml_program.global_store @global_mutable_undef = %0 : tensor<?xi32>
  42   ml_program.return
  43 }
  44
  45 // CHECK-LABEL: @global_load_store_tokens
  46 ml_program.subgraph @global_load_store_tokens() -> (tensor<?xi32>, !ml_program.token) {
  47   %token1 = ml_program.token
  48   %0, %token2 = ml_program.global_load_graph @global_mutable_undef
  49       ordering(() -> !ml_program.token) : tensor<?xi32>
  50   %token3 = ml_program.global_store_graph @global_mutable_undef = %0
  51       ordering(%token1, %token2 -> !ml_program.token) : tensor<?xi32>
  52   %token4 = ml_program.global_store_graph @global_mutable_undef = %0
  53       ordering(%token3 -> !ml_program.token) : tensor<?xi32>
  54
  55   ml_program.output %0, %token3 : tensor<?xi32>, !ml_program.token
  56 }