diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp --- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp +++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp @@ -2436,6 +2436,9 @@ resultElementTy.isInteger(8)) { Value index = rewriter.create(loc, rewriter.getIndexType(), inputValue); + Value offset = rewriter.create(loc, 128); + index = rewriter.create(loc, rewriter.getIndexType(), index, + offset); Value extract = rewriter.create(loc, table, ValueRange{index}); rewriter.create(loc, extract); diff --git a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir --- a/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir +++ b/mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir @@ -1101,17 +1101,21 @@ // ----- // CHECK-LABEL: @table8 -func @table8(%arg0: tensor<6xi8>, %arg1: tensor<513xi8>) -> () { +func @table8(%arg0: tensor<6xi8>, %arg1: tensor<512xi8>) -> () { // CHECK: %[[INIT:.+]] = linalg.init_tensor [6] // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi8>) outs(%[[INIT]] : tensor<6xi8>) // CHECK: ^bb0(%[[ARG_IN:.+]]: i8, %[[ARG_INIT:.+]]: i8) // CHECK: %[[CAST:.+]] = index_cast %[[ARG_IN]] - // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[CAST]]] + // CHECK: %[[OFFSET:.+]] = constant 128 + // CHECK: %[[ADD:.+]] = addi %[[CAST]], %[[OFFSET]] + // CHECK: %[[EXTRACT:.+]] = tensor.extract %arg1[%[[ADD]]] // CHECK: linalg.yield %[[EXTRACT]] - %0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<513xi8>) -> (tensor<6xi8>) + %0 = "tosa.table"(%arg0, %arg1) : (tensor<6xi8>, tensor<512xi8>) -> (tensor<6xi8>) return } +// ----- + // CHECK-LABEL: @table16 func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () { // CHECK: %[[INIT:.+]] = linalg.init_tensor [6]