[Mlir-commits] [mlir] [mlir][TensorToSPIRV] Add check for `tensor.extract` in TensorToSPIRV (PR #107110)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/107110

This patch add a check for `tensor.extract` in TensorToSPIRV, which is only support integers and floats type. Fix #74466.

>From 6fe4e1ffe9bc0b7ecdd73fdbda3e1b65b4a7caf6 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <moulongsheng at huawei.com>
Date: Tue, 3 Sep 2024 21:39:54 +0800
Subject: [PATCH] [mlir][TensorToSPIRV] Add check for `tensor.extract` in
 TensorToSPIRV

This patch add a check for `tensor.extract` in TensorToSPIRV, which
is only support integers and floats type.
---
 mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp    |  3 +++
 .../Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir  | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
index 0fb58623bdafbe..03bd79c843158d 100644
--- a/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
+++ b/mlir/lib/Conversion/TensorToSPIRV/TensorToSPIRV.cpp
@@ -45,6 +45,9 @@ class TensorExtractPattern final
                   ConversionPatternRewriter &rewriter) const override {
     auto tensorType = cast<RankedTensorType>(extractOp.getTensor().getType());
 
+    if (!tensorType.getElementType().isIntOrFloat())
+      return rewriter.notifyMatchFailure(extractOp,
+                                         "only integers and floats supported");
     if (!tensorType.hasStaticShape())
       return rewriter.notifyMatchFailure(extractOp, "non-static tensor");
 
diff --git a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
index 32d0fbea65b164..b1c2d5c2712eaa 100644
--- a/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/TensorToSPIRV/tensor-ops-to-spirv.mlir
@@ -29,6 +29,16 @@ func.func @tensor_extract_constant(%a : index, %b: index, %c: index) -> i32 {
 
 // -----
 
+// CHECK-LABEL: tensor_extract_unsupported_type
+func.func @tensor_extract_unsupported_type(%a : index) {
+  %cst = arith.constant dense<[1, 2]> : tensor<2xindex>
+  // CHECK: tensor.extract
+  %extract = tensor.extract %cst[%a] : tensor<2xindex>
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // Type conversion
 //===----------------------------------------------------------------------===//