[flang] add hlfir.transpose operation

[tblah]

Summary

Add a HLFIR operation for the TRANSPOSE transformational intrinsic,
according to the design set out in flang/doc/HighLevelFIR.md

[jeanPerier]

Small nit, looks good otherwise.

flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp

133

getFortranElementOrSequenceType(type).isa<fir::SequenceType>() should simplify the logic.

[tblah]

Thanks for review.

Changes: use isa<> instead of !!cast<>

flang/include/flang/Optimizer/HLFIR/HLFIRDialect.h

 /// Is this a fir.box or fir.class address or value type?
 inline bool isBoxAddressOrValueType(mlir::Type type) {
   return fir::unwrapRefType(type).isa<fir::BaseBoxType>();
 }
 
 bool isFortranScalarNumericalType(mlir::Type);
 bool isFortranNumericalArrayObject(mlir::Type);
 bool isFortranNumericalOrLogicalArrayObject(mlir::Type);
+bool isFortranArrayObject(mlir::Type);
 bool isPassByRefOrIntegerType(mlir::Type);
 bool isI1Type(mlir::Type);
 // scalar i1 or logical, or sequence of logical (via (boxed?) array or expr)
 bool isMaskArgument(mlir::Type);
 
 } // namespace hlfir
 
 #endif // FORTRAN_OPTIMIZER_HLFIR_HLFIRDIALECT_H

flang/include/flang/Optimizer/HLFIR/HLFIROpBase.td

 def AnyFortranNumericalArrayObject : Type<IsFortranNumericalArrayObjectPred,
     "any array-like object containing a numerical type">;
 
 def IsFortranNumericalOrLogicalArrayObjectPred
         : CPred<"::hlfir::isFortranNumericalOrLogicalArrayObject($_self)">;
 def AnyFortranNumericalOrLogicalArrayObject : Type<IsFortranNumericalOrLogicalArrayObjectPred,
     "any array-like object containing a numerical or logical type">;
 
+def IsFortranArrayObjectPred
+        : CPred<"::hlfir::isFortranArrayObject($_self)">;
+def AnyFortranArrayObject : Type<IsFortranArrayObjectPred,
+    "any array-like object">;
+
 def IsPassByRefOrIntegerTypePred
         : CPred<"::hlfir::isPassByRefOrIntegerType($_self)">;
 def AnyPassByRefOrIntegerType : Type<IsPassByRefOrIntegerTypePred,
     "an integer type either by value or by reference">;
 
 def IsMaskArgumentPred
         : CPred<"::hlfir::isMaskArgument($_self)">;
 def AnyFortranLogicalOrI1ArrayObject : Type<IsMaskArgumentPred,
     "A scalar i1 or logical or an array-like object containing logicals">;
 
 #endif // FORTRAN_DIALECT_HLFIR_OP_BASE

flang/include/flang/Optimizer/HLFIR/HLFIROps.td

def hlfir_MatmulOp : hlfir_Op<"matmul",
 
   let assemblyFormat = [{
     $lhs $rhs attr-dict `:` functional-type(operands, results)
   }];
 
   let hasVerifier = 1;
 }
 
+def hlfir_TransposeOp : hlfir_Op<"transpose", []> {
+  let summary = "TRANSPOSE transformational intrinsic";
+  let description = [{
+    Transpose a rank 2 array
+  }];
+
+  let arguments = (ins AnyFortranArrayObject:$array);
+
+  let results = (outs hlfir_ExprType);
+
+  let assemblyFormat = [{
+    $array attr-dict `:` functional-type(operands, results)
+  }];
+
+  let hasVerifier = 1;
+}
+
 def hlfir_AssociateOp : hlfir_Op<"associate", [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<fir_FortranVariableOpInterface>]> {
   let summary = "Create a variable from an expression value";
   let description = [{
     Create a variable from an expression value.
     For expressions, this operation is an incentive to re-use the expression
     storage, if any, after the bufferization pass when possible (if the
     expression is not used afterwards).

flang/lib/Optimizer/HLFIR/IR/HLFIRDialect.cpp

if (auto arrayTy =
           getFortranElementOrSequenceType(type).dyn_cast<fir::SequenceType>()) {
     mlir::Type eleTy = arrayTy.getEleTy();
     return isFortranScalarNumericalType(eleTy) ||
            mlir::isa<fir::LogicalType>(eleTy);
   }
   return false;
 }
 
+bool hlfir::isFortranArrayObject(mlir::Type type) {
+  if (isBoxAddressType(type))
+    return false;
+  return getFortranElementOrSequenceType(type).isa<fir::SequenceType>();
+}
+
 bool hlfir::isPassByRefOrIntegerType(mlir::Type type) {
   mlir::Type unwrappedType = fir::unwrapPassByRefType(type);
   return fir::isa_integer(unwrappedType);
 }
 
 bool hlfir::isI1Type(mlir::Type type) {
   if (mlir::IntegerType integer = type.dyn_cast<mlir::IntegerType>())
     if (integer.getWidth() == 1)

flang/lib/Optimizer/HLFIR/IR/HLFIROps.cpp

if (resultShape[0] != expectedResultShape[0])
     return emitOpError("incorrect result shape");
   if (resultShape.size() == 2 && resultShape[1] != expectedResultShape[1])
     return emitOpError("incorrect result shape");
 
   return mlir::success();
 }
 
 //===----------------------------------------------------------------------===//
+// TransposeOp
+//===----------------------------------------------------------------------===//
+
+mlir::LogicalResult hlfir::TransposeOp::verify() {
+  mlir::Value array = getArray();
+  fir::SequenceType arrayTy =
+      hlfir::getFortranElementOrSequenceType(array.getType())
+          .cast<fir::SequenceType>();
+  llvm::ArrayRef<int64_t> inShape = arrayTy.getShape();
+  std::size_t rank = inShape.size();
+  mlir::Type eleTy = arrayTy.getEleTy();
+  hlfir::ExprType resultTy = getResult().getType().cast<hlfir::ExprType>();
+  llvm::ArrayRef<int64_t> resultShape = resultTy.getShape();
+  std::size_t resultRank = resultShape.size();
+  mlir::Type resultEleTy = resultTy.getEleTy();
+
+  if (rank != 2 || resultRank != 2)
+    return emitOpError("input and output arrays should have rank 2");
+
+  if (inShape[0] != resultShape[1] || inShape[1] != resultShape[0])
+    return emitOpError("output shape does not match input array");
+
+  if (eleTy != resultEleTy)
+    return emitOpError(
+        "input and output arrays should have the same element type");
+
+  return mlir::success();
+}
+
+//===----------------------------------------------------------------------===//
 // AssociateOp
 //===----------------------------------------------------------------------===//
 
 void hlfir::AssociateOp::build(mlir::OpBuilder &builder,
                                mlir::OperationState &result, mlir::Value source,
                                llvm::StringRef uniq_name, mlir::Value shape,
                                mlir::ValueRange typeparams,
                                fir::FortranVariableFlagsAttr fortran_attrs) {

flang/test/HLFIR/invalid.fir

 // -----
 func.func @bad_matmul8(%arg0: !hlfir.expr<2xi32>, %arg1: !hlfir.expr<2x3xi32>) {
   // expected-error@+1 {{'hlfir.matmul' op incorrect result shape}}
   %0 = hlfir.matmul %arg0 %arg1 : (!hlfir.expr<2xi32>, !hlfir.expr<2x3xi32>) -> !hlfir.expr<1x3xi32>
   return
 }
 
 // -----
+func.func @bad_transpose1(%arg0: !hlfir.expr<2xi32>) {
+  // expected-error@+1 {{'hlfir.transpose' op input and output arrays should have rank 2}}
+  %0 = hlfir.transpose %arg0 : (!hlfir.expr<2xi32>) -> !hlfir.expr<2xi32>
+  return
+}
+
+// -----
+func.func @bad_transpose2(%arg0: !hlfir.expr<2x3xi32>) {
+  // expected-error@+1 {{'hlfir.transpose' op output shape does not match input array}}
+  %0 = hlfir.transpose %arg0 : (!hlfir.expr<2x3xi32>) -> !hlfir.expr<2x2xi32>
+  return
+}
+
+// -----
+func.func @bad_transpose3(%arg0: !hlfir.expr<2x3xi32>) {
+  // expected-error@+1 {{'hlfir.transpose' op input and output arrays should have the same element type}}
+  %0 = hlfir.transpose %arg0 : (!hlfir.expr<2x3xi32>) -> !hlfir.expr<3x2xf64>
+  return
+}
+
+// -----
 func.func @bad_assign_1(%arg0: !fir.box<!fir.array<?xi32>>, %arg1: !fir.box<!fir.array<?xi32>>) {
   // expected-error@+1 {{'hlfir.assign' op lhs must be an allocatable when `realloc` is set}}
   hlfir.assign %arg1 to %arg0 realloc : !fir.box<!fir.array<?xi32>>, !fir.box<!fir.array<?xi32>>
   return
 }
 
 // -----
 func.func @bad_assign_2(%arg0: !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>, %arg1: !fir.box<!fir.array<?xi32>>) {
   // expected-error@+1 {{'hlfir.assign' op `realloc` must be set and lhs must be a character allocatable when `keep_lhs_length_if_realloc` is set}}
   hlfir.assign %arg1 to %arg0 realloc keep_lhs_len : !fir.box<!fir.array<?xi32>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
   return
 }

flang/test/HLFIR/transpose.fir

+// Test hlfir.transpose operation parse, verify (no errors), and unparse
+
+// RUN: fir-opt %s | fir-opt | FileCheck %s
+
+// square matrix of known shape
+func.func @transpose0(%arg0: !hlfir.expr<2x2xi32>) {
+  %res = hlfir.transpose %arg0 : (!hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose0
+// CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x2xi32>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!hlfir.expr<2x2xi32>) -> !hlfir.expr<2x2xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+
+// rectangular matrix of known shape
+func.func @transpose1(%arg0: !hlfir.expr<2x3xi32>) {
+  %res = hlfir.transpose %arg0 : (!hlfir.expr<2x3xi32>) -> !hlfir.expr<3x2xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose1
+// CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x3xi32>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!hlfir.expr<2x3xi32>) -> !hlfir.expr<3x2xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+
+// matrix of assumed shape
+func.func @transpose2(%arg0: !hlfir.expr<?x?xi32>) {
+  %res = hlfir.transpose %arg0 : (!hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose2
+// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x?xi32>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!hlfir.expr<?x?xi32>) -> !hlfir.expr<?x?xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+
+// matrix where only some dimensions are known #1
+func.func @transpose3(%arg0: !hlfir.expr<?x2xi32>) {
+  %res = hlfir.transpose %arg0 : (!hlfir.expr<?x2xi32>) -> !hlfir.expr<2x?xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose3
+// CHECK:           %[[ARG0:.*]]: !hlfir.expr<?x2xi32>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!hlfir.expr<?x2xi32>) -> !hlfir.expr<2x?xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+
+// matrix where only some dimensions are known #2
+func.func @transpose4(%arg0: !hlfir.expr<2x?xi32>) {
+  %res = hlfir.transpose %arg0 : (!hlfir.expr<2x?xi32>) -> !hlfir.expr<?x2xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose4
+// CHECK:           %[[ARG0:.*]]: !hlfir.expr<2x?xi32>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!hlfir.expr<2x?xi32>) -> !hlfir.expr<?x2xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+
+// matrix is a boxed array
+func.func @transpose5(%arg0: !fir.box<!fir.array<1x2xi32>>) {
+  %res = hlfir.transpose %arg0 : (!fir.box<!fir.array<1x2xi32>>) -> !hlfir.expr<2x1xi32>
+  return
+}
+// CHECK-LABEL: func.func @transpose5
+// CHECK:           %[[ARG0:.*]]: !fir.box<!fir.array<1x2xi32>>
+// CHECK-NEXT:    %[[RES:.*]] = hlfir.transpose %[[ARG0]] : (!fir.box<!fir.array<1x2xi32>>) -> !hlfir.expr<2x1xi32>
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }