TosaToTensor: Support reshape on unsigned (#179)

Author: mgehre-amd

mlir/include/mlir/Conversion/TosaToTensor/TosaToTensor.h

@@ -16,6 +16,7 @@
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class TypeConverter;
 
 #define GEN_PASS_DECL_TOSATOTENSOR
 #include "mlir/Conversion/Passes.h.inc"
@@ -24,7 +25,8 @@ namespace tosa {
 
 std::unique_ptr<Pass> createTosaToTensor();
 
-void populateTosaToTensorConversionPatterns(RewritePatternSet *patterns);
+void populateTosaToTensorConversionPatterns(TypeConverter &converter,
+                                            RewritePatternSet *patterns);
 
 } // namespace tosa
 } // namespace mlir

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

@@ -207,7 +207,11 @@ class ReshapeConverterCollapseExpand
   matchAndRewrite(tosa::ReshapeOp reshape, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const final {
     ShapedType operandTy = cast<ShapedType>(adaptor.getInput1().getType());
-    ShapedType resultTy = cast<ShapedType>(reshape.getType());
+    ShapedType resultTy = cast_if_present<ShapedType>(getTypeConverter()->convertType(reshape.getType()));
+    if (!resultTy) {
+      return rewriter.notifyMatchFailure(
+        reshape.getLoc(), "could not convert result type");
+    }
     bool isDynamic = !operandTy.hasStaticShape();
 
     SmallVector<int64_t> intermediateShape;
@@ -218,7 +222,7 @@ class ReshapeConverterCollapseExpand
                    "the given two shapes");
     }
     auto intermediateTy = RankedTensorType::get(
-        intermediateShape, reshape.getType().getElementType());
+        intermediateShape, resultTy.getElementType());
 
     Value collapse = createCollapse(rewriter, reshape.getLoc(), intermediateTy,
                                     adaptor.getInput1());
@@ -415,9 +419,9 @@ struct ConcatConverter : public OpConversionPattern<tosa::ConcatOp> {
 } // namespace
 
 void mlir::tosa::populateTosaToTensorConversionPatterns(
-    RewritePatternSet *patterns) {
+    TypeConverter &converter, RewritePatternSet *patterns) {
   patterns->add<SliceConverter, PadConverter, ConcatConverter>(
       patterns->getContext());
 
-  patterns->add<ReshapeConverterCollapseExpand>(patterns->getContext());
+  patterns->add<ReshapeConverterCollapseExpand>(converter, patterns->getContext());
 }

mlir/lib/Conversion/TosaToTensor/TosaToTensorPass.cpp

@@ -12,6 +12,7 @@
 
 #include "mlir/Conversion/TosaToTensor/TosaToTensor.h"
 
+#include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
@@ -21,6 +22,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
+
 namespace mlir {
 #define GEN_PASS_DEF_TOSATOTENSOR
 #include "mlir/Conversion/Passes.h.inc"
@@ -42,7 +44,10 @@ struct TosaToTensor : public impl::TosaToTensorBase<TosaToTensor> {
     target.addLegalDialect<arith::ArithDialect>();
     target.addLegalDialect<tensor::TensorDialect>();
 
-    mlir::tosa::populateTosaToTensorConversionPatterns(&patterns);
+    TypeConverter converter;
+    mlir::tosa::populateTosaToLinalgTypeConversion(converter);
+
+    mlir::tosa::populateTosaToTensorConversionPatterns(converter, &patterns);
 
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))

mlir/test/Conversion/TosaToTensor/tosa-to-tensor.mlir

@@ -56,6 +56,19 @@ func.func @test_reshape_samerank(%arg0: tensor<3x2xf32>) -> tensor<2x3xf32> {
 
 // -----
 
+// CHECK-LABEL: @test_reshape_samerank_unsigned
+//  CHECK-SAME: (%[[ARG0:.*]]: tensor<3x2xui8>)
+func.func @test_reshape_samerank_unsigned(%arg0: tensor<3x2xui8>) -> tensor<2x3xui8> {
+  // CHECK-NEXT: %[[CAST1:.*]] = builtin.unrealized_conversion_cast %[[ARG0]] : tensor<3x2xui8> to tensor<3x2xi8>
+  // CHECK-NEXT: %[[RESHAPE1:.*]] = tensor.collapse_shape %[[CAST1]] {{\[}}[0, 1]] : tensor<3x2xi8> into tensor<6xi8>
+  // CHECK-NEXT: %[[RESHAPE2:.*]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1]] : tensor<6xi8> into tensor<2x3xi8>
+  // CHECK-NEXT: %[[CAST2:.*]] = builtin.unrealized_conversion_cast %[[RESHAPE2]] : tensor<2x3xi8> to tensor<2x3xui8>
+  %0 = "tosa.reshape"(%arg0) {new_shape = array<i64: 2, 3>} : (tensor<3x2xui8>) -> tensor<2x3xui8>
+  // CHECK-NEXT: return %[[CAST2]]
+  return %0 : tensor<2x3xui8>
+}
+// -----
+
 // CHECK-LABEL: @test_reshape_samerank_dyn
 //  CHECK-SAME: (%[[ARG0:.*]]: tensor<?x2xf32>)
 func.func @test_reshape_samerank_dyn(%arg0: tensor<?x2xf32>) -> tensor<2x?xf32> {