[Mlir-commits] [mlir] d325d2b - [mlir][tosa] Fix tosa::Select to linalg::generic indexingMaps bug

[Jacques Pienaar]

Author: AmosLewis
Date: 2022-10-10T15:47:50-07:00
New Revision: d325d2b407fcb7f37230b86fb291ef777497137d

URL: https://github.com/llvm/llvm-project/commit/d325d2b407fcb7f37230b86fb291ef777497137d
DIFF: https://github.com/llvm/llvm-project/commit/d325d2b407fcb7f37230b86fb291ef777497137d.diff

LOG: [mlir][tosa] Fix tosa::Select to linalg::generic indexingMaps bug

The tosa to linalg generic ops indexingMaps rank use is wrong.
	 Find this bug in gpt2 pytorch model lowering to tosa.
	 issue link is here https://github.com/llvm/llvm-project/issues/58154

Reviewed By: jpienaar

Differential Revision: https://reviews.llvm.org/D135343

Added: 
    

Modified: 
    mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
    mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 7aae079410989..63a491664052f 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -614,7 +614,7 @@ elementwiseMatchAndRewriteHelper(Operation *operation,
 
     operands.push_back(operand);
     indexingMaps.push_back(AffineMap::get(
-        /*dimCount=*/type.getRank(), /*symbolCount=*/0, affineExprs,
+        /*dimCount=*/rank, /*symbolCount=*/0, affineExprs,
         rewriter.getContext()));
   }
 

diff  --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
index a7297281d5ab3..a9b43d2b9f929 100644
--- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
+++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
@@ -1990,3 +1990,18 @@ func.func @resize_dyn(%input: tensor<?x2x2x1xi8>) -> () {
   %output = "tosa.resize"(%input) { scale = [4, 2, 4, 2], offset = [-1, -1], border = [1, 1], mode = "BILINEAR" } : (tensor<?x2x2x1xi8>)  -> (tensor<?x4x4x1xi32>)
   return
 }
+
+// -----
+
+// Regression test for using the wrong rank.
+
+// CHECK-DAG: affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK-DAG: affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
+// CHECK-DAG: affine_map<(d0, d1, d2, d3) -> ()>
+// CHECK-LABEL: @select_fp32
+func.func @select_fp32(%arg0: tensor<1x1x5x5xi1>, %arg1: tensor<1x12x5x5xf32>, %arg2: tensor<f32>) -> tensor<1x12x5x5xf32> {
+  // CHECK: linalg.generic
+  %0 = "tosa.select"(%arg0, %arg1, %arg2) : (tensor<1x1x5x5xi1>, tensor<1x12x5x5xf32>, tensor<f32>) -> tensor<1x12x5x5xf32>
+  return %0 : tensor<1x12x5x5xf32>
+}
+