mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir

   1 // RUN: mlir-opt --split-input-file --convert-tensor-to-spirv \
   2 // RUN:   --verify-diagnostics %s | FileCheck %s
   3
   4 //===----------------------------------------------------------------------===//
   5 // tensor.extract
   6 //===----------------------------------------------------------------------===//
   7
   8 // CHECK-LABEL: func @tensor_extract_constant
   9 // CHECK-SAME: (%[[A:.+]]: i32, %[[B:.+]]: i32, %[[C:.+]]: i32)
  10 func.func @tensor_extract_constant(%a : index, %b: index, %c: index) -> i32 {
  11   // CHECK: %[[CST:.+]] = spirv.Constant dense<[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]>
  12   %cst = arith.constant dense<[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]> : tensor<2x2x3xi32>
  13   // CHECK: %[[VAR:.+]] = spirv.Variable : !spirv.ptr<!spirv.array<12 x i32>, Function>
  14   // CHECK: spirv.Store "Function" %[[VAR]], %[[CST]] : !spirv.array<12 x i32>
  15   // CHECK: %[[C0:.+]] = spirv.Constant 0 : i32
  16   // CHECK: %[[C6:.+]] = spirv.Constant 6 : i32
  17   // CHECK: %[[MUL0:.+]] = spirv.IMul %[[A]], %[[C6]] : i32
  18   // CHECK: %[[C3:.+]] = spirv.Constant 3 : i32
  19   // CHECK: %[[MUL1:.+]] = spirv.IMul %[[B]], %[[C3]] : i32
  20   // CHECK: %[[ADD1:.+]] = spirv.IAdd %[[MUL1]], %[[MUL0]] : i32
  21   // CHECK: %[[C1:.+]] = spirv.Constant 1 : i32
  22   // CHECK: %[[ADD2:.+]] = spirv.IAdd %[[C]], %[[ADD1]] : i32
  23   // CHECK: %[[AC:.+]] = spirv.AccessChain %[[VAR]][%[[ADD2]]]
  24   // CHECK: %[[VAL:.+]] = spirv.Load "Function" %[[AC]] : i32
  25   %extract = tensor.extract %cst[%a, %b, %c] : tensor<2x2x3xi32>
  26   // CHECK: spirv.ReturnValue %[[VAL]]
  27   return %extract : i32
  28 }
  29
  30 // -----
  31
  32 // CHECK-LABEL: test_spirv_unsupported_type_index
  33 func.func @test_spirv_unsupported_type_index(%a : index) {
  34   %cst = arith.constant dense<[1, 2]> : tensor<2xindex>
  35   // CHECK: tensor.extract
  36   %extract = tensor.extract %cst[%a] : tensor<2xindex>
  37   return
  38 }
  39
  40 // CHECK-LABEL: test_spirv_unsupported_type_i128
  41 func.func @test_spirv_unsupported_type_i128(%a : index) {
  42   %cst = arith.constant dense<[1, 2]> : tensor<2xi128>
  43   // CHECK: tensor.extract
  44   %extract = tensor.extract %cst[%a] : tensor<2xi128>
  45   return
  46 }
  47
  48 // -----
  49
  50 //===----------------------------------------------------------------------===//
  51 // Type conversion
  52 //===----------------------------------------------------------------------===//
  53
  54 // CHECK-LABEL: func @tensor_0d
  55 // CHECK-NEXT:    spirv.Constant 1 : i32
  56 func.func @tensor_0d() -> () {
  57   %x = arith.constant dense<1> : tensor<i32>
  58   return
  59 }
  60
  61 // CHECK-LABEL: func @tensor_1d
  62 // CHECK-NEXT:    spirv.Constant dense<[1, 2, 3]> : tensor<3xi32> : !spirv.array<3 x i32>
  63 func.func @tensor_1d() -> () {
  64   %x = arith.constant dense<[1, 2, 3]> : tensor<3xi32>
  65   return
  66 }
  67
  68 // CHECK-LABEL: func @tensor_2d
  69 // CHECK-NEXT:    spirv.Constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32> : !spirv.array<6 x i32>
  70 func.func @tensor_2d() -> () {
  71   %x = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>
  72   return
  73 }
  74
  75 // We do not handle zero-element tensors yet. Just make we do not crash on them.
  76 // CHECK-LABEL: func @tensor_2d_empty
  77 // CHECK-NEXT:    arith.constant dense<>
  78 func.func @tensor_2d_empty() -> () {
  79   %x = arith.constant dense<> : tensor<2x0xi32>
  80   return
  81 }