[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

1028–1030

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

1036–1038

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

1038

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

1042

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

1055

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.

1056–1057

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

1072

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

Same for below

1115

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

1133

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

1167–1169

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

673–674

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

685

delete trailing blank lines

[rsuderman]

Updated for round 1 review

[rsuderman]

Missed comment

[rsuderman]

Fixed clang-tidy error

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

rewriter.replaceOpWithNewOp<linalg::GenericOp>(
         [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange args) {
           nestedBuilder.create<linalg::YieldOp>(op.getLoc(), *args.begin());
         });
 
     return success();
   }
 };
 
+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 reduction operation");
+
+    // First fill the output buffer with the init value.
+    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) {
+      rewriter.notifyMatchFailure(op,
+                                  "unable to create body for tosa.table op.");
+      return failure();
+    }
+
+    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::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>, ConcatConverter, ReshapeConverter,
-      RescaleConverter, ReverseConverter, TransposeConverter>(
+      RescaleConverter, ReverseConverter, TableConverter, TransposeConverter>(
       patterns->getContext());
 }

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

func @reverse(%arg0: tensor<5x4xi32>) -> () {
   // CHECK: [[INIT:%.+]] = linalg.init_tensor [5, 4]
   // CHECK: [[GENERIC:%.+]] = linalg.generic {indexing_maps = [#[[$MAP2]], #[[$MAP1]]], iterator_types = ["parallel", "parallel"]} ins(%arg0 : tensor<5x4xi32>) outs([[INIT]] : tensor<5x4xi32>) {
   // CHECK: ^bb0(%arg1: i32, %arg2: i32):
   // CHECK:   linalg.yield %arg1 : i32
   %1 = "tosa.reverse"(%arg0) {axis = 1 : i64} : (tensor<5x4xi32>) -> tensor<5x4xi32>
 
   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
+}
+