[mlir][tosa] Add tosa.table lowering to linalg.generic

[rsuderman]

Summary

Table op lowering to linalg.generic for both i8 (behaves like a gather) and a
pair of gathers with a quantized interpolation.

[rsuderman]

Fixed error message.

[rsuderman]

Removed superfluous comment.

[rsuderman]

Added class comment for Table Lowering.

[rsuderman]

Updated return statement.

[hanchung]

Sorry for the late review...

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

1422–1424

This is not a template class, so we don't need template here.

1430–1432

Can we move this above *ElementTy declarations? It would look more straight-forward to me.

1432

I think probably say require inputTy to have static shape? Meaning it is a requirement for matching.

1436

[optional] I personally would use ArrayRef<Value>{}

1449

I would rather create a block manually than use nested builder. Because

1. There are less indent.
2. You are creating a complicated body.
3. You can return a message when failing on creating the body.

I will use nested builder if it short.

If you are going to use nested body builder, I recently learned that there is ImplicitLocOpBuilder b(nextedLoc, nestedBuilder);. With the builder, you no longer need to pass the loc variable. It would look much simpler to me if you use ImplicitLocOpBuilder.

1450–1451

I think we can relax the condition to be "if all the element types are integer type and they have the same bitwidth?

1466

I feel this could be nestedLoc, nestedBuilder.getI32IntegerAttr(32768)

Same for below

1509

nit: there are two space right after +, delete one.

1527

drop punctuation at the send as these are supposed to be sentence fragments

1557–1559

not need to change in this patch. I'm the one who also prefer wrap clang-format on/off here, so we wont see this kind of change...

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

864–865

One tip (because I'm not sure if this is a convention). There was a discussion long time ago for capturing % in variables. We ended with no because we can avoid {{\[}}.

E.g., if you capture the variable without %, you can write

// CHECK: %[[IDX_CAST:.+]] = index_cast %[[IDX]]
...
// CHECK: %[[BASE:.+]] = tensor.extract %arg1[%[[IDX_CAST]]]

876

delete trailing blank lines

[rsuderman]

Updated for round 1 review

[rsuderman]

Missed comment

[rsuderman]

Fixed clang-tidy error

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

if (didEncounterError)
       return rewriter.notifyMatchFailure(
           argmaxOp, "unsupported tosa.argmax element type");
 
     rewriter.replaceOp(argmaxOp, linalgOp.getResult(0));
     return success();
   }
 };
 
+// Lowerings the TableOp to a series of gathers and numerica operations. This
+// includes interpolation between the high/low values. For the I8 varient, this
+// simplifies to a single gather operation.
+class TableConverter : public OpRewritePattern<tosa::TableOp> {
+public:
+  using OpRewritePattern<tosa::TableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::TableOp op,
+                                PatternRewriter &rewriter) const final {
+    auto loc = op.getLoc();
+    Value input = op.input();
+    Value table = op.table();
+    auto inputTy = input.getType().cast<ShapedType>();
+    auto tableTy = table.getType().cast<ShapedType>();
+    auto resultTy = op.getType().cast<ShapedType>();
+
+    if (!inputTy.hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          op, "require input type to have static shape");
+
+    auto inputElementTy = inputTy.getElementType();
+    auto tableElementTy = tableTy.getElementType();
+    auto resultElementTy = resultTy.getElementType();
+
+    auto initTensor =
+        rewriter
+            .create<linalg::InitTensorOp>(loc, ArrayRef<Value>{},
+                                          resultTy.getShape(), resultElementTy)
+            .result();
+
+    SmallVector<AffineMap, 2> affineMaps = {
+        rewriter.getMultiDimIdentityMap(resultTy.getRank()),
+        rewriter.getMultiDimIdentityMap(resultTy.getRank())};
+
+    auto genericOp = rewriter.create<linalg::GenericOp>(
+        loc, resultTy, ValueRange({input}), ValueRange{initTensor}, affineMaps,
+        getNParallelLoopsAttrs(resultTy.getRank()));
+    rewriter.replaceOp(op, genericOp.getResult(0));
+
+    {
+      OpBuilder::InsertionGuard regionGuard(rewriter);
+      Block *block =
+          rewriter.createBlock(&genericOp.region(), genericOp.region().end(),
+                               TypeRange({inputElementTy, resultElementTy}));
+
+      auto inputValue = block->getArgument(0);
+      rewriter.setInsertionPointToStart(block);
+      if (inputElementTy.isInteger(8) && tableElementTy.isInteger(8) &&
+          resultElementTy.isInteger(8)) {
+        Value index = rewriter.create<IndexCastOp>(loc, rewriter.getIndexType(),
+                                                   inputValue);
+        Value extract =
+            rewriter.create<tensor::ExtractOp>(loc, table, ValueRange{index});
+        rewriter.create<linalg::YieldOp>(loc, extract);
+        return success();
+      }
+
+      if (inputElementTy.isInteger(16) && tableElementTy.isInteger(16) &&
+          resultElementTy.isInteger(32)) {
+        Value extend = rewriter.create<SignExtendIOp>(
+            loc, rewriter.getI32Type(), inputValue);
+
+        auto offset =
+            rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(32768));
+        auto seven =
+            rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(7));
+        auto one =
+            rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(1));
+        auto b1111111 =
+            rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(127));
+
+        // Compute the index and fractional part from the input value:
+        // value = value + 32768
+        // index = value >> 7;
+        // fraction = 0x01111111 & value
+        auto extendAdd = rewriter.create<AddIOp>(loc, extend, offset);
+        Value index =
+            rewriter.create<UnsignedShiftRightOp>(loc, extendAdd, seven);
+        Value fraction = rewriter.create<mlir::AndOp>(loc, extendAdd, b1111111);
+
+        // Extract the base and next values from the table.
+        // base = (int32_t) table[index];
+        // next = (int32_t) table[index + 1];
+        Value indexPlusOne = rewriter.create<AddIOp>(loc, index, one);
+
+        index =
+            rewriter.create<IndexCastOp>(loc, rewriter.getIndexType(), index);
+        indexPlusOne = rewriter.create<IndexCastOp>(
+            loc, rewriter.getIndexType(), indexPlusOne);
+
+        Value base =
+            rewriter.create<tensor::ExtractOp>(loc, table, ValueRange{index});
+        Value next = rewriter.create<tensor::ExtractOp>(
+            loc, table, ValueRange{indexPlusOne});
+
+        base = rewriter.create<SignExtendIOp>(loc, rewriter.getI32Type(), base);
+        next = rewriter.create<SignExtendIOp>(loc, rewriter.getI32Type(), next);
+
+        // Use the fractional part to interpolate between the input values:
+        // result = (base << 7) + (next - base) * fraction
+        Value baseScaled = rewriter.create<ShiftLeftOp>(loc, base, seven);
+        Value diff = rewriter.create<SubIOp>(loc, next, base);
+        Value diffScaled = rewriter.create<MulIOp>(loc, diff, fraction);
+        Value result = rewriter.create<AddIOp>(loc, baseScaled, diffScaled);
+
+        rewriter.create<linalg::YieldOp>(loc, result);
+
+        return success();
+      }
+    }
+
+    return rewriter.notifyMatchFailure(
+        op, "unable to create body for tosa.table op");
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<
-      PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
-      PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::ReciprocalOp>,
-      PointwiseConverter<tosa::NegateOp>, PointwiseConverter<tosa::PowOp>,
-      PointwiseConverter<tosa::RsqrtOp>, PointwiseConverter<tosa::LogOp>,
-      PointwiseConverter<tosa::ExpOp>, PointwiseConverter<tosa::AbsOp>,
-      PointwiseConverter<tosa::TanhOp>, PointwiseConverter<tosa::BitwiseAndOp>,
+      // clang-format off
+      PointwiseConverter<tosa::AddOp>,
+      PointwiseConverter<tosa::SubOp>,
+      PointwiseConverter<tosa::MulOp>,
+      PointwiseConverter<tosa::NegateOp>,
+      PointwiseConverter<tosa::PowOp>,
+      PointwiseConverter<tosa::ReciprocalOp>,
+      PointwiseConverter<tosa::RsqrtOp>,
+      PointwiseConverter<tosa::LogOp>,
+      PointwiseConverter<tosa::ExpOp>,
+      PointwiseConverter<tosa::AbsOp>,
+      PointwiseConverter<tosa::TanhOp>,
+      PointwiseConverter<tosa::BitwiseAndOp>,
       PointwiseConverter<tosa::BitwiseOrOp>,
       PointwiseConverter<tosa::BitwiseNotOp>,
       PointwiseConverter<tosa::BitwiseXorOp>,
       PointwiseConverter<tosa::LogicalAndOp>,
       PointwiseConverter<tosa::LogicalNotOp>,
       PointwiseConverter<tosa::LogicalOrOp>,
-      PointwiseConverter<tosa::LogicalXorOp>, PointwiseConverter<tosa::CastOp>,
+      PointwiseConverter<tosa::LogicalXorOp>,
+      PointwiseConverter<tosa::CastOp>,
       PointwiseConverter<tosa::LogicalLeftShiftOp>,
       PointwiseConverter<tosa::LogicalRightShiftOp>,
-      PointwiseConverter<tosa::SelectOp>, PointwiseConverter<tosa::GreaterOp>,
+      PointwiseConverter<tosa::SelectOp>,
+      PointwiseConverter<tosa::GreaterOp>,
       PointwiseConverter<tosa::GreaterEqualOp>,
-      PointwiseConverter<tosa::MaximumOp>, PointwiseConverter<tosa::MinimumOp>,
-      PointwiseConverter<tosa::CeilOp>, PointwiseConverter<tosa::FloorOp>,
-      PointwiseConverter<tosa::ClampOp>, PointwiseConverter<tosa::ReluNOp>,
-      PointwiseConverter<tosa::SigmoidOp>, IdentityNConverter<tosa::IdentityOp>,
-      IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceAllOp>,
-      ReduceConverter<tosa::ReduceAnyOp>, ReduceConverter<tosa::ReduceMinOp>,
-      ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
-      ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter,
-      PadConverter, ReshapeConverter, RescaleConverter, ReverseConverter,
-      TileConverter, TransposeConverter, MatMulConverter,
+      PointwiseConverter<tosa::MaximumOp>,
+      PointwiseConverter<tosa::MinimumOp>,
+      PointwiseConverter<tosa::CeilOp>,
+      PointwiseConverter<tosa::FloorOp>,
+      PointwiseConverter<tosa::ClampOp>,
+      PointwiseConverter<tosa::ReluNOp>,
+      PointwiseConverter<tosa::SigmoidOp>,
+      IdentityNConverter<tosa::IdentityOp>,
+      IdentityNConverter<tosa::IdentityNOp>,
+      ReduceConverter<tosa::ReduceAllOp>,
+      ReduceConverter<tosa::ReduceAnyOp>,
+      ReduceConverter<tosa::ReduceMinOp>,
+      ReduceConverter<tosa::ReduceMaxOp>,
+      ReduceConverter<tosa::ReduceSumOp>,
+      ReduceConverter<tosa::ReduceProdOp>,
+      ArgMaxConverter,
+      ConcatConverter,
+      PadConverter,
+      ReshapeConverter,
+      RescaleConverter,
+      ReverseConverter,
+      TableConverter,
+      TileConverter,
+      TransposeConverter,
+      MatMulConverter,
       FullyConnectedConverter>(patterns->getContext());
+      // clang-format on
 }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @argmax(%arg0 : tensor<3x2xi32>, %arg1 : tensor<6xf32>) -> () {
   // CHECK: cmpf ogt
   // CHECK: select
   // CHECK: select
   // CHECK: linalg.yield
   %2 = "tosa.argmax"(%arg1) { axis = 0 : i64} : (tensor<6xf32>)  -> (tensor<i32>)
 
   return
 }
+
+// -----
+
+// CHECK-LABEL: @table8
+func @table8(%arg0: tensor<6xi8>, %arg1: tensor<513xi8>) -> () {
+  // 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:   linalg.yield %[[EXTRACT]]
+  %0 = "tosa.table"(%arg0, %arg1)  : (tensor<6xi8>, tensor<513xi8>)  -> (tensor<6xi8>)
+  return
+}
+
+// CHECK-LABEL: @table16
+func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () {
+  // CHECK: %[[INIT:.+]] = linalg.init_tensor [6]
+  // CHECK: %[[GENERIC:.+]] = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg0 : tensor<6xi16>) outs(%[[INIT]] : tensor<6xi32>)
+  // CHECK: ^bb0(%arg2: i16, %arg3: i32)
+  // CHECK: %[[EXT_IN:.+]] = sexti %arg2
+  // CHECK: %[[C32768:.+]] = constant 32768
+  // CHECK: %[[C7:.+]] = constant 7
+  // CHECK: %[[C1:.+]] = constant 1
+  // CHECK: %[[C127:.+]] = constant 127
+  // CHECK: %[[INADD:.+]] = addi %[[EXT_IN]], %[[C32768]]
+  // CHECK: %[[IDX:.+]] = shift_right_unsigned %[[INADD]], %[[C7]]
+  // CHECK: %[[FRACTION:.+]] = and %[[INADD]], %[[C127]]
+  // CHECK: %[[IDXPLUS1:.+]] = addi %[[IDX]], %[[C1]]
+  // CHECK: %[[IDX_CAST:.+]] = index_cast %[[IDX]]
+  // CHECK: %[[IDXPLUS1_CAST:.+]] = index_cast %[[IDXPLUS1]]
+  // CHECK: %[[BASE:.+]] = tensor.extract %arg1[%[[IDX_CAST]]]
+  // CHECK: %[[NEXT:.+]] = tensor.extract %arg1[%[[IDXPLUS1_CAST]]]
+  // CHECK: %[[BASE_EXT:.+]] = sexti %[[BASE]]
+  // CHECK: %[[NEXT_EXT:.+]] = sexti %[[NEXT]]
+  // CHECK: %[[BASE_MUL:.+]] = shift_left %[[BASE_EXT]], %[[C7]]
+  // CHECK: %[[DIFF:.+]] = subi %[[NEXT_EXT]], %[[BASE_EXT]]
+  // CHECK: %[[DIFF_MUL:.+]] = muli %[[DIFF]], %[[FRACTION]]
+  // CHECK: %[[RESULT:.+]] = addi %[[BASE_MUL]], %[[DIFF_MUL]]
+  // CHECK: linalg.yield %[[RESULT]]
+  %0 = "tosa.table"(%arg0, %arg1)  : (tensor<6xi16>, tensor<513xi16>)  -> (tensor<6xi32>)
+  return
+}