82 lines · plain
1// RUN: mlir-opt --split-input-file --convert-tensor-to-spirv \2// RUN: --verify-diagnostics %s | FileCheck %s3 4//===----------------------------------------------------------------------===//5// tensor.extract6//===----------------------------------------------------------------------===//7 8// CHECK-LABEL: func @tensor_extract_constant9// CHECK-SAME: (%[[A:.+]]: i32, %[[B:.+]]: i32, %[[C:.+]]: i32)10func.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 : i3216 // CHECK: %[[C6:.+]] = spirv.Constant 6 : i3217 // CHECK: %[[MUL0:.+]] = spirv.IMul %[[A]], %[[C6]] : i3218 // CHECK: %[[C3:.+]] = spirv.Constant 3 : i3219 // CHECK: %[[MUL1:.+]] = spirv.IMul %[[B]], %[[C3]] : i3220 // CHECK: %[[ADD1:.+]] = spirv.IAdd %[[MUL1]], %[[MUL0]] : i3221 // CHECK: %[[C1:.+]] = spirv.Constant 1 : i3222 // CHECK: %[[ADD2:.+]] = spirv.IAdd %[[C]], %[[ADD1]] : i3223 // CHECK: %[[AC:.+]] = spirv.AccessChain %[[VAR]][%[[ADD2]]]24 // CHECK: %[[VAL:.+]] = spirv.Load "Function" %[[AC]] : i3225 %extract = tensor.extract %cst[%a, %b, %c] : tensor<2x2x3xi32>26 // CHECK: spirv.ReturnValue %[[VAL]]27 return %extract : i3228}29 30// -----31 32// CHECK-LABEL: test_spirv_unsupported_type_index33func.func @test_spirv_unsupported_type_index(%a : index) {34 %cst = arith.constant dense<[1, 2]> : tensor<2xindex>35 // CHECK: tensor.extract36 %extract = tensor.extract %cst[%a] : tensor<2xindex>37 return38}39 40// CHECK-LABEL: test_spirv_unsupported_type_i12841func.func @test_spirv_unsupported_type_i128(%a : index) {42 %cst = arith.constant dense<[1, 2]> : tensor<2xi128>43 // CHECK: tensor.extract44 %extract = tensor.extract %cst[%a] : tensor<2xi128>45 return46}47 48// -----49 50//===----------------------------------------------------------------------===//51// Type conversion52//===----------------------------------------------------------------------===//53 54// CHECK-LABEL: func @tensor_0d55// CHECK-NEXT: spirv.Constant 1 : i3256func.func @tensor_0d() -> () {57 %x = arith.constant dense<1> : tensor<i32>58 return59}60 61// CHECK-LABEL: func @tensor_1d62// CHECK-NEXT: spirv.Constant dense<[1, 2, 3]> : tensor<3xi32> : !spirv.array<3 x i32>63func.func @tensor_1d() -> () {64 %x = arith.constant dense<[1, 2, 3]> : tensor<3xi32>65 return66}67 68// CHECK-LABEL: func @tensor_2d69// CHECK-NEXT: spirv.Constant dense<[1, 2, 3, 4, 5, 6]> : tensor<6xi32> : !spirv.array<6 x i32>70func.func @tensor_2d() -> () {71 %x = arith.constant dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi32>72 return73}74 75// We do not handle zero-element tensors yet. Just make we do not crash on them.76// CHECK-LABEL: func @tensor_2d_empty77// CHECK-NEXT: arith.constant dense<>78func.func @tensor_2d_empty() -> () {79 %x = arith.constant dense<> : tensor<2x0xi32>80 return81}82