[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

1391–1393

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

1399–1401

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

1401

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

1405

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

1418

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.

1419–1420

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

1435

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

Same for below

1478

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

1496

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

1531–1533

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

827–828

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]]]

839

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().template cast<ShapedType>();
+    auto tableTy = table.getType().template cast<ShapedType>();
+    auto resultTy = op.getType().template cast<ShapedType>();
+
+    auto inputElementTy = inputTy.getElementType();
+    auto tableElementTy = tableTy.getElementType();
+    auto resultElementTy = resultTy.getElementType();
+
+    if (!inputTy.hasStaticShape())
+      return rewriter.notifyMatchFailure(
+          op, "No initial value found for table operation");
+
+    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())};
+
+    bool bodySucceeded = true;
+
+    rewriter.replaceOpWithNewOp<linalg::GenericOp>(
+        op, resultTy, ValueRange({input}), ValueRange{initTensor}, affineMaps,
+        getNParallelLoopsAttrs(resultTy.getRank()),
+        [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
+          if (inputElementTy.isInteger(8) && tableElementTy.isInteger(8) &&
+              resultElementTy.isInteger(8)) {
+            Value index = nestedBuilder.create<IndexCastOp>(
+                nestedLoc, rewriter.getIndexType(), args[0]);
+            Value extract = nestedBuilder.create<tensor::ExtractOp>(
+                nestedLoc, table, ValueRange{index});
+            nestedBuilder.create<linalg::YieldOp>(nestedLoc, extract);
+            return;
+          }
+
+          if (inputElementTy.isInteger(16) && tableElementTy.isInteger(16) &&
+              resultElementTy.isInteger(32)) {
+            Value extend = rewriter.create<SignExtendIOp>(
+                nestedLoc, rewriter.getI32Type(), args[0]);
+
+            auto offset = nestedBuilder.create<ConstantIntOp>(
+                nestedLoc, 32768, nestedBuilder.getI32Type());
+
+            auto seven = nestedBuilder.create<ConstantIntOp>(
+                nestedLoc, 7, nestedBuilder.getI32Type());
+
+            auto one = nestedBuilder.create<ConstantIntOp>(
+                nestedLoc, 1, nestedBuilder.getI32Type());
+            auto b1111111 = nestedBuilder.create<ConstantIntOp>(
+                nestedLoc, 127, nestedBuilder.getI32Type());
+
+            // Compute the index and fractional part from the input value:
+            // value = value + 32768
+            // index = value >> 7;
+            // fraction = 0x01111111 & value
+            auto extendAdd =
+                nestedBuilder.create<AddIOp>(nestedLoc, extend, offset);
+            Value index = nestedBuilder.create<UnsignedShiftRightOp>(
+                nestedLoc, extendAdd, seven);
+            Value fraction = nestedBuilder.create<mlir::AndOp>(
+                nestedLoc, extendAdd, b1111111);
+
+            // Extract the base and next values from the table.
+            // base = (int32_t) table[index];
+            // next = (int32_t) table[index + 1];
+            Value indexPlusOne =
+                nestedBuilder.create<AddIOp>(nestedLoc, index, one);
+
+            index = nestedBuilder.create<IndexCastOp>(
+                nestedLoc, nestedBuilder.getIndexType(), index);
+            indexPlusOne = nestedBuilder.create<IndexCastOp>(
+                nestedLoc, nestedBuilder.getIndexType(), indexPlusOne);
+
+            Value base = nestedBuilder.create<tensor::ExtractOp>(
+                nestedLoc, table, ValueRange{index});
+            Value next = nestedBuilder.create<tensor::ExtractOp>(
+                nestedLoc, table, ValueRange{indexPlusOne});
+
+            base = nestedBuilder.create<SignExtendIOp>(
+                nestedLoc, nestedBuilder.getI32Type(), base);
+            next = nestedBuilder.create<SignExtendIOp>(
+                nestedLoc, nestedBuilder.getI32Type(), next);
+
+            // Use the fractional part to interpolate between the input values:
+            // result = (base << 7) +  (next - base) * fraction
+            Value baseScaled =
+                nestedBuilder.create<ShiftLeftOp>(nestedLoc, base, seven);
+            Value diff = nestedBuilder.create<SubIOp>(nestedLoc, next, base);
+            Value diffScaled =
+                nestedBuilder.create<MulIOp>(nestedLoc, diff, fraction);
+            Value result =
+                nestedBuilder.create<AddIOp>(nestedLoc, baseScaled, diffScaled);
+
+            nestedBuilder.create<linalg::YieldOp>(nestedLoc, result);
+            return;
+          }
+
+          bodySucceeded = false;
+        });
+
+    if (!bodySucceeded)
+      return rewriter.notifyMatchFailure(
+          op, "unable to create body for tosa.table op.");
+
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<
       PointwiseConverter<tosa::AddOp>, PointwiseConverter<tosa::SubOp>,
       PointwiseConverter<tosa::MulOp>, PointwiseConverter<tosa::NegateOp>,
       PointwiseConverter<tosa::PowOp>, PointwiseConverter<tosa::RsqrtOp>,
patterns->add<
       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>,
       IdentityNConverter<tosa::IdentityOp>,
       IdentityNConverter<tosa::IdentityNOp>, ReduceConverter<tosa::ReduceMinOp>,
       ReduceConverter<tosa::ReduceMaxOp>, ReduceConverter<tosa::ReduceSumOp>,
-      ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter, PadConverter,
-      ReshapeConverter, RescaleConverter, ReverseConverter, TileConverter,
-      TransposeConverter, MatMulConverter, FullyConnectedConverter>(
-        patterns->getContext());
+      ReduceConverter<tosa::ReduceProdOp>, ArgMaxConverter, ConcatConverter,
+      PadConverter, ReshapeConverter, RescaleConverter, ReverseConverter,
+      TableConverter, TileConverter, TransposeConverter, MatMulConverter,
+      FullyConnectedConverter>(patterns->getContext());
 }

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(%0 : 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
+}
+