[MLIR] Shape to Standard dialect lowering

[frgossen]

Summary

Add new pass to lower operations from the shape to the std dialect.
The conversion applies only to the size_to_index and to the index_to_size
operations.
Patterns will be added as needed.

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

22

nit: no need for the empty line

23

this class is not really needed. You can just move addLegalDialect<scf::SCFDialect, StandardOpsDialect>(); together with addLegalOp below.

32

nit: no need for the empty line

35

nit: no need for the empty line

50

nit: no need for the empty line

53

and here

64

and even here

74

nit: and here

79

maybe just inline it in addConversion?

81

why do you need to initialize it with nullptr?

100

merge this with the first addLegalOp call ?

121

you have two patterns here, but one test.

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

10

I don't think this comment brings a lot. Maybe smth like:

// CHECK-LABEL: @size_to_index
func @size_to_index(%size : !shape.size) -> !shape.size {
  return %size : !shape.size
}
// CHECK-SAME: ([[SIZE:%.*]]: index)
// CHECK: return [[SIZE]] : index

[frgossen]

Fix nits and add test cases

[frgossen]

Thanks

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

81

It's not really needed here, rather a habit that prevents uninitialised pointers in general.
I'm happy to make it uninitialised here.

100

Good eye, thanks

121

It's actually worse, I had no test for any of them

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

10

Updated the test cases and comments. Hope this makes it clear.

[pifon2a]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

62

since you're copying ctx ptr into the lambda, you probably don't need it here.

[frgossen]

Remove unused member

[rriddle]

mlir/include/mlir/Conversion/ShapeToStandard/ShapeToStandard.h

12

Can you trim this header? It doesn't seem necessary.

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

17

using namespace mlir;

93

This isn't llvm style, please use explicit resolution: mlir:: populateShapeToStandardConversionPatterns

Same below.

[frgossen]

Addressed in https://reviews.llvm.org/D81160

mlir/include/mlir/Conversion/Passes.td

 }
 
 def ConvertParallelLoopToGpu : Pass<"convert-parallel-loops-to-gpu"> {
   let summary = "Convert mapped scf.parallel ops to gpu launch operations";
   let constructor = "mlir::createParallelLoopToGpuPass()";
 }
 
 //===----------------------------------------------------------------------===//
+// ShapeToStandard
+//===----------------------------------------------------------------------===//
+
+def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
+  let summary = "Convert operations from the shape dialect into the standard "
+                "dialect";
+  let constructor = "mlir::createConvertShapeToStandardPass()";
+}
+
+//===----------------------------------------------------------------------===//
 // StandardToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertStandardToLLVM : Pass<"convert-std-to-llvm", "ModuleOp"> {
   let summary = "Convert scalar and vector operations from the Standard to the "
                 "LLVM dialect";
   let description = [{
     Convert standard operations into the LLVM IR dialect operations.

mlir/include/mlir/Conversion/ShapeToStandard/ShapeToStandard.h

+//===- ShapeToStandard.h - Conversion utils from shape to std 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_SHAPETOSTANDARD_SHAPETOSTANDARD_H_
+#define MLIR_CONVERSION_SHAPETOSTANDARD_SHAPETOSTANDARD_H_
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+
+class MLIRContext;
+class ModuleOp;
+template <typename T>
+class OperationPass;
+
+void populateShapeToStandardConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx);
+
+std::unique_ptr<OperationPass<ModuleOp>> createConvertShapeToStandardPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_SHAPETOSTANDARD_SHAPETOSTANDARD_H_

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/GPUToSPIRV/ConvertGPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/LinalgToLLVM/LinalgToLLVM.h"
 #include "mlir/Conversion/LinalgToSPIRV/LinalgToSPIRVPass.h"
 #include "mlir/Conversion/LinalgToStandard/LinalgToStandard.h"
 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
+#include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/Affine/Passes.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/LegalizeForExport.h"
 #include "mlir/Dialect/Linalg/Passes.h"

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(AffineToStandard)
 add_subdirectory(AVX512ToLLVM)
 add_subdirectory(GPUCommon)
 add_subdirectory(GPUToNVVM)
 add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(GPUToVulkan)
 add_subdirectory(LinalgToLLVM)
 add_subdirectory(LinalgToSPIRV)
 add_subdirectory(LinalgToStandard)
 add_subdirectory(SCFToGPU)
 add_subdirectory(SCFToStandard)
+add_subdirectory(ShapeToStandard)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)

mlir/lib/Conversion/ShapeToStandard/CMakeLists.txt

+add_mlir_conversion_library(MLIRShapeToStandard
+  ShapeToStandard.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/ShapeToStandard
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_COMPONENTS
+  Core
+
+  LINK_LIBS PUBLIC
+  MLIREDSC
+  MLIRIR
+  MLIRShape
+  MLIRPass
+  MLIRSCF
+  MLIRTransforms
+  )

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

+//===- LinalgToStandard.cpp - conversion from Linalg to Standard 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/ShapeToStandard/ShapeToStandard.h"
+
+#include "../PassDetail.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Shape/IR/Shape.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace {
+
+/// Conversion patterns.
+class SizeToIndexOpConversion
+    : public OpConversionPattern<shape::SizeToIndexOp> {
+public:
+  using OpConversionPattern<shape::SizeToIndexOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::SizeToIndexOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::SizeToIndexOpOperandAdaptor transformed(operands);
+    rewriter.replaceOp(op.getOperation(), transformed.arg());
+    return success();
+  }
+};
+
+class IndexToSizeOpConversion
+    : public OpConversionPattern<shape::IndexToSizeOp> {
+public:
+  using OpConversionPattern<shape::IndexToSizeOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(shape::IndexToSizeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    shape::IndexToSizeOpOperandAdaptor transformed(operands);
+    rewriter.replaceOp(op.getOperation(), transformed.arg());
+    return success();
+  }
+};
+
+/// Type conversions.
+class ShapeTypeConverter : public TypeConverter {
+public:
+  using TypeConverter::convertType;
+
+  ShapeTypeConverter(MLIRContext *ctx) {
+    // Add default pass-through conversion.
+    addConversion([&](Type type) { return type; });
+    addConversion([ctx](shape::SizeType type) { return IndexType::get(ctx); });
+  }
+};
+
+/// Conversion pass.
+class ConvertShapeToStandardPass
+    : public ConvertShapeToStandardBase<ConvertShapeToStandardPass> {
+
+  void runOnOperation() override {
+
+    // Setup type conversion.
+    MLIRContext &ctx = getContext();
+    ShapeTypeConverter typeConverter(&ctx);
+
+    // Setup target legality.
+    ConversionTarget target(ctx);
+    target.addLegalDialect<scf::SCFDialect, StandardOpsDialect>();
+    target.addLegalOp<ModuleOp, ModuleTerminatorOp, ReturnOp>();
+    target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+      return typeConverter.isSignatureLegal(op.getType());
+    });
+
+    // Setup conversion patterns.
+    OwningRewritePatternList patterns;
+    populateShapeToStandardConversionPatterns(patterns, &ctx);
+    populateFuncOpTypeConversionPattern(patterns, &ctx, typeConverter);
+
+    // Apply conversion.
+    auto module = getOperation();
+    if (failed(applyFullConversion(module, target, patterns, &typeConverter)))
+      signalPassFailure();
+  }
+};
+
+} // namespace
+
+void populateShapeToStandardConversionPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *ctx) {
+  // clang-format off
+  patterns.insert<
+      IndexToSizeOpConversion,
+      SizeToIndexOpConversion>(ctx);
+  // clang-format on
+}
+
+std::unique_ptr<OperationPass<ModuleOp>> createConvertShapeToStandardPass() {
+  return std::make_unique<ConvertShapeToStandardPass>();
+}
+
+} // namespace mlir

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

+// RUN: mlir-opt --split-input-file --convert-shape-to-std --verify-diagnostics %s | FileCheck %s --dump-input-on-failure
+
+// Convert `size` to `index` type.
+// CHECK-LABEL: @size_id
+// CHECK-SAME: (%[[SIZE:.*]]: index)
+func @size_id(%size : !shape.size) -> !shape.size {
+  // CHECK: return %[[SIZE]] : index
+  return %size : !shape.size
+}
+
+// -----
+
+// Lower `size_to_index` conversion to no-op.
+// CHECK-LABEL: @size_to_index
+// CHECK-SAME: (%[[SIZE:.*]]: index) -> index
+func @size_to_index(%size : !shape.size) -> index {
+  // CHECK-NEXT: return %[[SIZE]] : index
+  %index = shape.size_to_index %size
+  return %index : index
+}
+
+// -----
+
+// Lower `index_to_size` conversion to no-op.
+// CHECK-LABEL: @index_to_size
+// CHECK-SAME: (%[[INDEX:.*]]: index) -> index
+func @index_to_size(%index : index) -> !shape.size {
+  // CHECK-NEXT: return %[[INDEX]] : index
+  %size = shape.index_to_size %index
+  return %size : !shape.size
+}