[mlir] Add pass to convert elementwise ops to linalg.

[silvas]

Summary

This patch converts elementwise ops on tensors to linalg.generic ops
with the same elementwise op in the payload (except rewritten to
operate on scalars, obviously). This is a great form for later fusion to
clean up.

E.g.

// Compute: %arg0 + %arg1 - %arg2
func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
  %0 = addf %arg0, %arg1 : tensor<?xf32>
  %1 = subf %0, %arg2 : tensor<?xf32>
  return %1 : tensor<?xf32>
}

Running this through
mlir-opt -convert-std-to-linalg -linalg-fusion-for-tensor-ops we get:

func @f(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>, %arg2: tensor<?xf32>) -> tensor<?xf32> {
  %0 = linalg.generic {indexing_maps = [#map0, #map0, #map0, #map0], iterator_types = ["parallel"]} ins(%arg0, %arg1, %arg2 : tensor<?xf32>, tensor<?xf32>, tensor<?xf32>) {
  ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):  // no predecessors
    %1 = addf %arg3, %arg4 : f32
    %2 = subf %1, %arg5 : f32
    linalg.yield %2 : f32 
  } -> tensor<?xf32>
  return %0 : tensor<?xf32>
}

So the elementwise ops on tensors have nicely collapsed into a single
linalg.generic, which is the form we want for further transformations.

[mehdi_amini]

mlir/include/mlir/Conversion/Passes.td

360

Can you add a more comprehensive "description" field?

mlir/lib/Conversion/PassDetail.h

38

Closing comment? (clang-tidy)

[nicolasvasilache]

LG, thanks Sean!
Approving conditioned on the doc improvement.
The elementwise trait is also something we want for vectorization so it will be doubly helpful.

[silvas]

address comments

[silvas]

Update to use new elementwise trait

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/OpenMPToLLVM/ConvertOpenMPToLLVM.h"
 #include "mlir/Conversion/PDLToPDLInterp/PDLToPDLInterp.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/ConvertSPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Conversion/StandardToLinalg/StandardToLinalg.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/VectorToSPIRV/ConvertVectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

let options = [
     Option<"dataLayout", "data-layout", "std::string",
            /*default=*/"\"\"",
            "String description (LLVM format) of the data layout that is "
            "expected on the produced module">
   ];
 }
 
 //===----------------------------------------------------------------------===//
+// StandardToLinalg
+//===----------------------------------------------------------------------===//
+
+def ConvertStandardToLinalg : Pass<"convert-std-to-linalg", "FuncOp"> {
+  let summary = "Convert std ops to linalg";
+  let description = [{
+    Convert ops from the std dialect to linalg.
+
+    Currently, this pass only supports std elementwise ops operating on tensors
+    and converts them to linalg parallel loops.
+  }];
+  let constructor = "mlir::createConvertStandardToLinalgPass()";
+  let dependentDialects = ["linalg::LinalgDialect"];
+}
+
+//===----------------------------------------------------------------------===//
 // StandardToSPIRV
 //===----------------------------------------------------------------------===//
 
 def LegalizeStandardForSPIRV : Pass<"legalize-std-for-spirv"> {
   let summary = "Legalize standard ops for SPIR-V lowering";
   let constructor = "mlir::createLegalizeStdOpsForSPIRVLoweringPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
 }

mlir/include/mlir/Conversion/StandardToLinalg/StandardToLinalg.h

+//===- StandardToLinalg.h - Conversion utils from std to linalg -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_STANDARDTOLINALG_STANDARDTOLINALG_H_
+#define MLIR_CONVERSION_STANDARDTOLINALG_STANDARDTOLINALG_H_
+
+#include <memory>
+
+namespace mlir {
+
+class Pass;
+class MLIRContext;
+class Pass;
+class OwningRewritePatternList;
+
+void populateStandardToLinalgConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
+std::unique_ptr<Pass> createConvertStandardToLinalgPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_STANDARDTOLINALG_STANDARDTOLINALG_H_

mlir/integration_test/Dialect/Linalg/CPU/test-std-elementwise.mlir

+// RUN: mlir-opt %s -convert-std-to-linalg -linalg-bufferize -convert-linalg-to-loops -convert-linalg-to-llvm -convert-std-to-llvm | \
+// RUN: mlir-cpu-runner -e main -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_runner_utils%shlibext \
+// RUN: | FileCheck %s
+
+func @main() {
+  %a = constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
+  %b = constant dense<[10.0, 20.0, 30.0]> : tensor<3xf32>
+
+  %addf = addf %a, %b : tensor<3xf32>
+  %addf_unranked = tensor_cast %addf : tensor<3xf32> to tensor<*xf32>
+  call @print_memref_f32(%addf_unranked) : (tensor<*xf32>) -> ()
+  // CHECK: Unranked Memref base@ = {{.*}} rank = 1 offset = 0 sizes = [3] strides = [1] data =
+  // CHECK-NEXT: [11,  22,  33]
+
+  return
+}
+
+func @print_memref_f32(%ptr : tensor<*xf32>)

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(OpenMPToLLVM)
 add_subdirectory(PDLToPDLInterp)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
+add_subdirectory(StandardToLinalg)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
 add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/PassDetail.h

 class LLVMDialect;
 class LLVMAVX512Dialect;
 } // end namespace LLVM
 
 namespace NVVM {
 class NVVMDialect;
 } // end namespace NVVM
 
+namespace linalg {
+class LinalgDialect;
+} // end namespace linalg
+
 namespace pdl_interp {
 class PDLInterpDialect;
 } // end namespace pdl_interp
 
 namespace ROCDL {
 class ROCDLDialect;
 } // end namespace ROCDL
 

mlir/lib/Conversion/StandardToLinalg/CMakeLists.txt

+add_mlir_conversion_library(MLIRStandardToLinalg
+  StandardToLinalg.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/StandardToLinalg
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRLinalg
+  MLIRTransforms
+  )

mlir/lib/Conversion/StandardToLinalg/StandardToLinalg.cpp

+//===- StandardToLinalg.cpp - conversion from Std to Linalg dialect -------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/StandardToLinalg/StandardToLinalg.h"
+
+#include "../PassDetail.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+
+static bool isStdElementwiseOpOnRankedTensors(Operation *op) {
+  if (!llvm::all_of(op->getResultTypes(),
+                    [](Type type) { return type.isa<RankedTensorType>(); })) {
+    return false;
+  }
+  if (!llvm::all_of(op->getOperandTypes(),
+                    [](Type type) { return type.isa<RankedTensorType>(); })) {
+    return false;
+  }
+  // TODO: Have an "elementwise" trait that supercedes this.
+  // clang-format off
+  return isa<
+    AbsFOp,
+    AddFOp,
+    AddIOp,
+    AndOp,
+    AtanOp,
+    Atan2Op,
+    CeilFOp,
+    FloorFOp,
+    CmpFOp,
+    CmpIOp,
+    CopySignOp,
+    CosOp,
+    SinOp,
+    DivFOp,
+    ExpOp,
+    Exp2Op,
+    FPExtOp,
+    FPToSIOp,
+    FPToUIOp,
+    FPTruncOp,
+    IndexCastOp,
+    LogOp,
+    Log10Op,
+    Log2Op,
+    MulFOp,
+    MulIOp,
+    NegFOp,
+    OrOp,
+    RemFOp,
+    RsqrtOp,
+    SelectOp,
+    ShiftLeftOp,
+    SignedDivIOp,
+    SignedRemIOp,
+    SignedShiftRightOp,
+    SignExtendIOp,
+    SIToFPOp,
+    SqrtOp,
+    SubFOp,
+    SubIOp,
+    SubViewOp,
+    TanhOp,
+    TruncateIOp,
+    UIToFPOp,
+    UnsignedDivIOp,
+    UnsignedRemIOp,
+    UnsignedShiftRightOp,
+    XOrOp,
+    ZeroExtendIOp
+  >(op);
+  // clang-format on
+}
+
+namespace {
+struct ConvertStdElementwiseOpOnRankedTensors : public RewritePattern {
+  ConvertStdElementwiseOpOnRankedTensors()
+      : RewritePattern(/*benefit=*/1, MatchAnyOpTypeTag()) {}
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const final {
+    if (!isStdElementwiseOpOnRankedTensors(op))
+      return rewriter.notifyMatchFailure(op, "requires std elementwise op");
+
+    auto rank = op->getResult(0).getType().cast<RankedTensorType>().getRank();
+    SmallVector<AffineMap, 3> indexingMaps(
+        op->getNumResults() + op->getNumOperands(),
+        rewriter.getMultiDimIdentityMap(rank));
+    SmallVector<StringRef, 6> iteratorTypes(rank,
+                                            getParallelIteratorTypeName());
+    rewriter.replaceOpWithNewOp<linalg::GenericOp>(
+        op, /*resultTensorTypes=*/op->getResultTypes(),
+        /*inputs=*/op->getOperands(),
+        /*outputBuffers=*/ValueRange(),
+        /*initTensors=*/ValueRange(),
+        /*indexingMaps=*/indexingMaps,
+        /*iteratorTypes=*/iteratorTypes,
+        /*bodyBuilder=*/
+        [&](OpBuilder &builder, Location loc, ValueRange regionArgs) {
+          OperationState state(loc, op->getName());
+          state.addAttributes(op->getAttrs());
+          state.addOperands(regionArgs);
+          auto resultTypes = llvm::to_vector<6>(
+              llvm::map_range(op->getResultTypes(), [](Type type) {
+                return type.cast<TensorType>().getElementType();
+              }));
+          state.addTypes(resultTypes);
+          auto *scalarOp = builder.createOperation(state);
+          builder.create<linalg::YieldOp>(loc, scalarOp->getResults());
+        });
+    return success();
+  }
+};
+} // namespace
+
+void mlir::populateStandardToLinalgConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *) {
+  patterns.insert<ConvertStdElementwiseOpOnRankedTensors>();
+}
+
+namespace {
+class ConvertStandardToLinalgPass
+    : public ConvertStandardToLinalgBase<ConvertStandardToLinalgPass> {
+
+  void runOnOperation() final {
+    auto func = getOperation();
+    auto *context = &getContext();
+    ConversionTarget target(*context);
+    OwningRewritePatternList patterns;
+
+    populateStandardToLinalgConversionPatterns(patterns, context);
+    target.markUnknownOpDynamicallyLegal(
+        [](Operation *op) { return !isStdElementwiseOpOnRankedTensors(op); });
+
+    if (failed(applyPartialConversion(func, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createConvertStandardToLinalgPass() {
+  return std::make_unique<ConvertStandardToLinalgPass>();
+}

mlir/test/Conversion/StandardToLinalg/std-to-linalg.mlir

+// RUN: mlir-opt -convert-std-to-linalg -split-input-file %s | FileCheck %s
+
+// In-depth checking of the linalg.generic op for a very trivial case.
+// CHECK-LABEL: #map0 = affine_map<() -> ()>
+// CHECK-LABEL:   func @addf_rank0
+func @addf_rank0(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
+  // CHECK: %{{.*}} = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []} ins(%{{.*}}, %{{.*}} : tensor<f32>, tensor<f32>) {
+  // CHECK: ^bb0(%[[LHS:.*]]: f32, %[[RHS:.*]]: f32):
+  // CHECK:   %[[YIELD:.*]] = addf %[[LHS]], %[[RHS]] : f32
+  // CHECK:   linalg.yield %[[YIELD]] : f32
+  // CHECK: } -> tensor<f32>
+  %0 = addf %arg0, %arg1 : tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
+// Check indexing maps and iterator types for the rank > 0 case.
+// CHECK-LABEL: #map0 = affine_map<(d0) -> (d0)>
+// CHECK-LABEL: func @addf_rank1
+func @addf_rank1(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK: linalg.generic{{.*}}indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel"]
+  %0 = addf %arg0, %arg1 : tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+// Check a unary op.
+// CHECK-LABEL: func @exp
+func @exp(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: linalg.generic
+  // CHECK: ^bb0(%[[SCALAR:.*]]: f32):
+  // CHECK:   %[[YIELD:.*]] = exp %[[SCALAR]] : f32
+  // CHECK:   linalg.yield %[[YIELD]] : f32
+  %0 = exp %arg0 : tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
+// Check a case with varying operand types.
+// CHECK-LABEL: func @select
+func @select(%arg0: tensor<i1>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<i32> {
+  // CHECK: linalg.generic
+  // CHECK: ^bb0(%[[PRED:.*]]: i1, %[[TRUE_VAL:.*]]: i32, %[[FALSE_VAL:.*]]: i32):
+  // CHECK:   select %[[PRED]], %[[TRUE_VAL]], %[[FALSE_VAL]] : i32
+  %0 = select %arg0, %arg1, %arg2 : tensor<i1>, tensor<i32>
+  return %0 : tensor<i32>
+}
+
+// -----
+
+// Spot-check an op that requires copying attributes properly to the created scalar op.
+// CHECK-LABEL: func @cmpf(
+func @cmpf(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
+  // CHECK: cmpf "olt", %{{.*}}, %{{.*}} : f32
+  %0 = cmpf "olt", %arg0, %arg1 : tensor<f32>
+  return %0 : tensor<i1>
+}