[FXML-4791] Lower memref expand/collapse to EmitC

mlir/include/mlir/Dialect/EmitC/IR/EmitC.h

@@ -47,6 +47,9 @@ bool isSupportedFloatType(mlir::Type type);
 /// Determines whether \p type is a emitc.size_t/ssize_t type.
 bool isPointerWideType(mlir::Type type);
 
+/// Give the name of the EmitC reference attribute.
+StringRef getReferenceAttributeName();
+
 } // namespace emitc
 } // namespace mlir
 

mlir/include/mlir/Dialect/EmitC/IR/EmitC.td

@@ -265,8 +265,7 @@ def EmitC_CallOpaqueOp : EmitC_Op<"call_opaque", [CExpression]> {
 
 def EmitC_CastOp : EmitC_Op<"cast",
     [CExpression,
-     DeclareOpInterfaceMethods<CastOpInterface>,
-     SameOperandsAndResultShape]> {
+     DeclareOpInterfaceMethods<CastOpInterface>]> {
   let summary = "Cast operation";
   let description = [{
     The `cast` operation performs an explicit type conversion and is emitted
@@ -285,9 +284,11 @@ def EmitC_CastOp : EmitC_Op<"cast",
     ```
   }];
 
-  let arguments = (ins EmitCType:$source);
+  let arguments = (ins EmitCType:$source,
+                       UnitAttr:$reference);
   let results = (outs EmitCType:$dest);
-  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest) (`ref` $reference^)?";
+  let hasVerifier = 1;
 }
 
 def EmitC_CmpOp : EmitC_BinaryOp<"cmp", [CExpression]> {
@@ -1050,14 +1051,16 @@ def EmitC_GlobalOp : EmitC_Op<"global", [Symbol]> {
                        OptionalAttr<EmitC_OpaqueOrTypedAttr>:$initial_value,
                        UnitAttr:$extern_specifier,
                        UnitAttr:$static_specifier,
-                       UnitAttr:$const_specifier);
+                       UnitAttr:$const_specifier,
+                       UnitAttr:$reference);
 
   let assemblyFormat = [{
        (`extern` $extern_specifier^)?
        (`static` $static_specifier^)?
        (`const` $const_specifier^)?
        $sym_name
        `:` custom<EmitCGlobalOpTypeAndInitialValue>($type, $initial_value)
+       (`ref` $reference^)?
        attr-dict
   }];
 

mlir/lib/Conversion/MemRefToEmitC/MemRefToEmitC.cpp

@@ -15,6 +15,8 @@
 #include "mlir/Dialect/EmitC/IR/EmitC.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 
@@ -166,6 +168,71 @@ struct ConvertStore final : public OpConversionPattern<memref::StoreOp> {
     return success();
   }
 };
+
+struct ConvertCollapseShape final
+    : public OpConversionPattern<memref::CollapseShapeOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::CollapseShapeOp op, OpAdaptor operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto arrayValue = dyn_cast<TypedValue<emitc::ArrayType>>(operands.getSrc());
+    if (!arrayValue) {
+      return rewriter.notifyMatchFailure(op.getLoc(), "expected array type");
+    }
+
+    auto resultTy = getTypeConverter()->convertType(op.getType());
+    if (!resultTy) {
+      return rewriter.notifyMatchFailure(op.getLoc(),
+                                         "cannot convert result type");
+    }
+
+    // Do not generate casts between arrays with dynamic shapes
+    auto shape = arrayValue.getType().getShape();
+    for (auto d : shape) {
+      if (d == ShapedType::kDynamic)
+        return failure();
+    }
+    auto newCastOp = rewriter.create<emitc::CastOp>(op->getLoc(), resultTy,
+                                                    operands.getSrc());
+    newCastOp.setReference(true);
+    rewriter.replaceOp(op, newCastOp);
+    return success();
+  }
+};
+
+struct ConvertExpandShape final
+    : public OpConversionPattern<memref::ExpandShapeOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ExpandShapeOp op, OpAdaptor operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto arrayValue = dyn_cast<TypedValue<emitc::ArrayType>>(operands.getSrc());
+    if (!arrayValue) {
+      return rewriter.notifyMatchFailure(op.getLoc(), "expected array type");
+    }
+
+    auto resultTy = getTypeConverter()->convertType(op.getType());
+    if (!resultTy) {
+      return rewriter.notifyMatchFailure(op.getLoc(),
+                                         "cannot convert result type");
+    }
+
+    // Do not generate casts between arrays with dynamic shapes
+    auto shape = arrayValue.getType().getShape();
+    for (auto d : shape) {
+      if (d == ShapedType::kDynamic)
+        return failure();
+    }
+    auto newCastOp = rewriter.create<emitc::CastOp>(op->getLoc(), resultTy,
+                                                    operands.getSrc());
+    newCastOp.setReference(true);
+    rewriter.replaceOp(op, newCastOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void mlir::populateMemRefToEmitCTypeConversion(TypeConverter &typeConverter) {
@@ -187,5 +254,6 @@ void mlir::populateMemRefToEmitCTypeConversion(TypeConverter &typeConverter) {
 void mlir::populateMemRefToEmitCConversionPatterns(RewritePatternSet &patterns,
                                                    TypeConverter &converter) {
   patterns.add<ConvertAlloca, ConvertGlobal, ConvertGetGlobal, ConvertLoad,
-               ConvertStore>(converter, patterns.getContext());
+               ConvertStore, ConvertCollapseShape, ConvertExpandShape>(
+      converter, patterns.getContext());
 }

mlir/lib/Dialect/EmitC/IR/EmitC.cpp

@@ -10,6 +10,7 @@
 #include "mlir/Dialect/EmitC/IR/EmitCTraits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/IRMapping.h"
@@ -122,6 +123,8 @@ bool mlir::emitc::isPointerWideType(Type type) {
       type);
 }
 
+StringRef mlir::emitc::getReferenceAttributeName() { return "emitc.reference"; }
+
 /// Check that the type of the initial value is compatible with the operations
 /// result type.
 static LogicalResult verifyInitializationAttribute(Operation *op,
@@ -225,13 +228,39 @@ LogicalResult emitc::AssignOp::verify() {
 bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
   Type input = inputs.front(), output = outputs.front();
 
+  // Arrays & pointers don't cast to/from scalars
+  if (isa<emitc::ArrayType, emitc::PointerType>(input) !=
+      isa<emitc::ArrayType, emitc::PointerType>(output))
+    return false;
+
+  // Arrays & pointers can be casted to each other.
+  if (isa<emitc::ArrayType, emitc::PointerType>(input) &&
+      isa<emitc::ArrayType, emitc::PointerType>(output))
+    return true;
+
+  // Scalars
   return (
       (emitc::isIntegerIndexOrOpaqueType(input) ||
        emitc::isSupportedFloatType(input) || isa<emitc::PointerType>(input)) &&
       (emitc::isIntegerIndexOrOpaqueType(output) ||
        emitc::isSupportedFloatType(output) || isa<emitc::PointerType>(output)));
 }
 
+LogicalResult CastOp::verify() {
+  bool isReference = this->getReference();
+
+  if (isa<emitc::ArrayType>(getDest().getType())) {
+    if (!isReference)
+      return emitOpError("cast of array must bear a reference");
+    return success();
+  }
+
+  if (isReference)
+    return emitOpError("cast of value type must not bear a reference");
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // CallOp
 //===----------------------------------------------------------------------===//
@@ -945,6 +974,10 @@ LogicalResult emitc::ArrayType::verify(
   for (int64_t dim : shape) {
     if (dim < 0)
       return emitError() << "dimensions must have non-negative size";
+    // TODO support dynamic dimensions on the outermost dimension,
+    // like int (*a)[50][100]
+    if (dim == ShapedType::kDynamic)
+      return emitError() << "dimensions must have static size";
   }
 
   if (!elementType)