[mlir] Add structural type conversions for SCF dialect.

[silvas]

Summary

A "structural" type conversion is one where the underlying ops are
completely agnostic to the actual types involved, and simply need to
update their types accordingly.

An example of this is scf.if -- the values yielded by the regions merely
need to match the result types. The scf.if op itself just needs to
update its types accordingly with how the body is modified.

See the comment in TestBufferizeStructuralOps.cpp for more details on
the issues with the current formulation here. Any help on resolving
those issues would be greatly appreciated.

[herhut]

Thanks, looks good! Please wait for other reviewers to comment.

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

34

Should ResultRange and friends maybe grow a setTypes method?

43

nit: why not have one continuous namespace?

mlir/test/lib/Transforms/TestBufferizeStructuralOps.cpp

26

Could this be done by the infra itself based on RegionOpBranchInterface and ReturnLikeOp? Those are sufficient to understand region data flow, so the infra would know which operands of terminators need materialization.

[silvas]

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

43

LLVM coding standards suggest keeping "privatizing" annotations like static on functions or anoymous namespaces on classes scoped as tightly as possible: https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

mlir/test/lib/Transforms/TestBufferizeStructuralOps.cpp

26

I would love that. I'll be looking more into how to solve this soon.

[rriddle]

mlir/include/mlir/Dialect/SCF/Transforms.h

16

Drive by: This header looks like it could be replaced by forward declarations.

[silvas]

mlir/include/mlir/Dialect/SCF/Transforms.h

16

Do we have a documented forward declaration policy? I find them more error-prone so I avoid them if possible.

[rriddle]

mlir/include/mlir/Dialect/SCF/Transforms.h

16

https://llvm.org/docs/CodingStandards.html#include-as-little-as-possible

[silvas]

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

34

Good idea! I did it in a separate follow-on patch for clarity: https://reviews.llvm.org/D89591

[silvas]

Use forward decls

[mehdi_amini]

How does it work with legality: someone has to leave scf.for/scf.if illegal for the conversion to kick-in here? But we also want them legal after the conversion.
Seems like we'd want this kind of pattern to also come with their notion of legality: a "structural op" should be marked dynamically legal based on if its invariants are correctly preserved?

mlir/include/mlir/Dialect/SCF/Transforms.h

48

Can you document the public API?

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

21

Doc here as well?

28

s/auto/Type/

43

Here you have two classes that could be in the same anonymous namespace declaration I believe?

[silvas]

update

[silvas]

Abandoning this. See the improved patch in D89526

I finally figured out how to do the proper source/target type materializations to eliminate the need for a single mega-conversion pass, even for structural conversions :)

[mehdi_amini]

In D89526#2340594, @silvas wrote:

Abandoning this. See the improved patch in D89526

You're referencing the current revision here, error in copy/paste?

[mehdi_amini]

I think you meant to link to D89757?

[silvas]

In D89526#2340755, @mehdi_amini wrote:

I think you meant to link to D89757?

Correct

mlir/include/mlir/Dialect/SCF/Transforms.h

 //===- Transforms.h - SCF dialect transformation utilities ------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines transformations on SCF operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_TRANSFORMS_H_
 #define MLIR_DIALECT_SCF_TRANSFORMS_H_
 
 #include "llvm/ADT/ArrayRef.h"
 
 namespace mlir {
 
+class MLIRContext;
+class OwningRewritePatternList;
 class Region;
+class TypeConverter;
 
 namespace scf {
 
 class ParallelOp;
 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency
 /// analysis.
 void naivelyFuseParallelOps(Region &region);
 
 /// Tile a parallel loop of the form
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4, %arg5)
 ///
 /// into
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4*tileSize[0],
 ///                                                   %arg5*tileSize[1])
 ///     scf.parallel (%j0, %j1) = (0, 0) to (min(tileSize[0], %arg2-%j0)
 ///                                           min(tileSize[1], %arg3-%j1))
 ///                                        step (%arg4, %arg5)
 /// The old loop is replaced with the new one.
 void tileParallelLoop(ParallelOp op, llvm::ArrayRef<int64_t> tileSizes);
 
+void populateSCFStructuralTypeConversions(MLIRContext *context,
+                                          TypeConverter &typeConverter,
+                                          OwningRewritePatternList &patterns);
+
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_TRANSFORMS_H_

mlir/lib/Dialect/SCF/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransforms
   LoopSpecialization.cpp
   ParallelLoopFusion.cpp
   ParallelLoopTiling.cpp
+  StructuralTypeConversions.cpp
   Utils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF
 
   DEPENDS
   MLIRSCFPassIncGen
 

mlir/lib/Dialect/SCF/Transforms/StructuralTypeConversions.cpp

+//===- StructuralTypeConversions.cpp - scf structural type conversions ----===//
+//
+// 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 "PassDetail.h"
+#include "mlir/Dialect/SCF/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+using namespace mlir;
+using namespace mlir::scf;
+
+namespace {
+class ConvertForOpTypes : public OpConversionPattern<ForOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(ForOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Type, 6> newResultTypes;
+    for (auto type : op.getResultTypes()) {
+      Type newType = typeConverter->convertType(type);
+      if (!newType)
+        return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion");
+      newResultTypes.push_back(newType);
+    }
+    rewriter.updateRootInPlace(op, [&] {
+      for (auto t : llvm::zip(op.getResults(), newResultTypes))
+        std::get<0>(t).setType(std::get<1>(t));
+      auto bodyArgs = op.getBody()->getArguments();
+      for (auto t : llvm::zip(llvm::drop_begin(bodyArgs, 1), newResultTypes))
+        std::get<0>(t).setType(std::get<1>(t));
+    });
+    return success();
+  }
+};
+} // namespace
+
+namespace {
+class ConvertIfOpTypes : public OpConversionPattern<IfOp> {
+public:
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(IfOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    SmallVector<Type, 6> newResultTypes;
+    for (auto type : op.getResultTypes()) {
+      Type newType = typeConverter->convertType(type);
+      if (!newType)
+        return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion");
+      newResultTypes.push_back(newType);
+    }
+    rewriter.updateRootInPlace(op, [&] {
+      for (auto t : llvm::zip(op.getResults(), newResultTypes))
+        std::get<0>(t).setType(std::get<1>(t));
+    });
+    return success();
+  }
+};
+} // namespace
+
+void mlir::scf::populateSCFStructuralTypeConversions(
+    MLIRContext *context, TypeConverter &typeConverter,
+    OwningRewritePatternList &patterns) {
+  patterns.insert<ConvertIfOpTypes>(typeConverter, context);
+  patterns.insert<ConvertForOpTypes>(typeConverter, context);
+}

mlir/test/Transforms/bufferize-structural-ops.mlir

+// RUN: mlir-opt -test-bufferize-structural-ops -split-input-file %s | FileCheck %s
+
+// Basic cases.
+
+// CHECK-LABEL: func @identity(%arg0: memref<?xf32>) -> memref<?xf32> {
+// CHECK-NEXT:    return %arg0 : memref<?xf32>
+func @identity(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  return %arg0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @bb_arg(%arg0: memref<?xf32>) -> memref<?xf32> {
+// CHECK-NEXT:    br ^bb1(%arg0 : memref<?xf32>)
+// CHECK-NEXT:  ^bb1(%[[BBARG:.*]]: memref<?xf32>):
+// CHECK-NEXT:    return %[[BBARG]] : memref<?xf32>
+func @bb_arg(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+  br ^bb1(%arg0: tensor<?xf32>)
+^bb1(%bbarg: tensor<?xf32>):
+  return %bbarg : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @if(%arg0: i1, %arg1: memref<?xf32>, %arg2: memref<?xf32>) -> memref<?xf32> {
+// CHECK-NEXT:    %[[RET:.*]] = scf.if %arg0 -> (memref<?xf32>) {
+// CHECK-NEXT:      scf.yield %arg1 : memref<?xf32>
+// CHECK-NEXT:    } else {
+// CHECK-NEXT:      scf.yield %arg2 : memref<?xf32>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return %[[RET]] : memref<?xf32>
+func @if(%pred: i1, %true_val: tensor<?xf32>, %false_val: tensor<?xf32>) -> tensor<?xf32> {
+  %0 = scf.if %pred -> (tensor<?xf32>) {
+    scf.yield %true_val : tensor<?xf32>
+  } else {
+    scf.yield %false_val : tensor<?xf32>
+  }
+  return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: func @for(%arg0: memref<f32>, %arg1: index, %arg2: index, %arg3: index) -> memref<f32> {
+// CHECK-NEXT:    %[[RET:.*]] = scf.for %arg4 = %arg1 to %arg2 step %arg3 iter_args(%arg5 = %arg0) -> (memref<f32>) {
+// CHECK-NEXT:      scf.yield %arg5 : memref<f32>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    return %[[RET]] : memref<f32>
+// CHECK-NEXT:  }
+func @for(%arg0: tensor<f32>, %lb: index, %ub: index, %step: index) -> tensor<f32> {
+  %ret = scf.for %iv = %lb to %ub step %step iter_args(%iter = %arg0) -> tensor<f32> {
+    scf.yield %iter : tensor<f32>
+  }
+  return %ret : tensor<f32>
+}

mlir/test/lib/Transforms/CMakeLists.txt

 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestTransforms
   TestAllReduceLowering.cpp
   TestAffineLoopParametricTiling.cpp
   TestBufferPlacement.cpp
+  TestBufferizeStructuralOps.cpp
   TestExpandTanh.cpp
   TestCallGraph.cpp
   TestConstantFold.cpp
   TestConvVectorization.cpp
   TestConvertCallOp.cpp
   TestConvertGPUKernelToCubin.cpp
   TestConvertGPUKernelToHsaco.cpp
   TestDominance.cpp

mlir/test/lib/Transforms/TestBufferizeStructuralOps.cpp

+//===- TestBufferizeStructuralOps.cpp ---------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A "structural" type conversion is one where the underlying ops are completely
+// agnostic to the actual types involved, and simply need to update their types
+// accordingly.
+//
+// An example of this is scf.if -- the values yielded by the regions merely
+// need to match the result types. The scf.if op itself just needs to update
+// its types accordingly with how the body is modified.
+//
+// Due to a limitation of the dialect conversion infrastructure, all structural
+// conversions need to be done as a single mega-conversion-pass that does a full
+// conversion of types. The reason is that conversion patterns apply to the
+// results of the region-bearing op itself -- but the operands of yield-like
+// terminators inside the regions are not updated, causing verifier errors.
+// There doesn't seem to be a way to trigger the dialect conversion framework to
+// correctly update them.
+//
+// TODO: Should every structural type conversion just open-code updating its
+// yield-like terminators by calling
+// `getTypeConverter()->materializeTargetConversion(...)`?
+//
+// That is why we have a single standalone test pass instead of composable
+// individual (non-test) passes like we do for non-structural bufferization
+// patterns (which are composable because source/target materializations are
+// automatically inserted by the dialect conversion framework).
+//
+// Pragmatically, we use bufferization type conversion since it has a
+// well-established type conversion infrastructure in-tree, and this pass is
+// useful documentation for how downstream users can utilize this functionality,
+// given the limitations described above.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/SCF/Transforms.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Bufferize.h"
+
+using namespace mlir;
+
+namespace {
+struct TestBufferizeStructuralOpsPass
+    : mlir::PassWrapper<TestBufferizeStructuralOpsPass,
+                        OperationPass<ModuleOp>> {
+
+  void runOnOperation() override {
+    auto func = getOperation();
+    auto *context = &getContext();
+
+    BufferizeTypeConverter typeConverter;
+
+    OwningRewritePatternList patterns;
+
+    ConversionTarget target(*context);
+
+    // All ops whose results are not tensor types are legal.
+    target.markUnknownOpDynamicallyLegal([](Operation *op) {
+      return llvm::all_of(op->getResultTypes(),
+                          [](Type type) { return !type.isa<TensorType>(); });
+    });
+
+    populateFuncOpTypeConversionPattern(patterns, context, typeConverter);
+    target.addDynamicallyLegalOp<mlir::FuncOp>([&](mlir::FuncOp op) {
+      return typeConverter.isSignatureLegal(op.getType()) &&
+             typeConverter.isLegal(&op.getBody());
+    });
+
+    scf::populateSCFStructuralTypeConversions(context, typeConverter, patterns);
+
+    if (failed(applyFullConversion(func, target, patterns)))
+      return signalPassFailure();
+  };
+};
+} // end anonymous namespace
+
+namespace mlir {
+void registerTestBufferizeStructuralOpsPass() {
+  PassRegistration<TestBufferizeStructuralOpsPass>(
+      "test-bufferize-structural-ops",
+      "Tests the bufferization of 'structural ops' that cannot be composed "
+      "into individual passes");
+}
+} // end namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerSimpleParametricTilingPass();
 void registerSliceAnalysisTestPass();
 void registerSymbolTestPasses();
 void registerTestAffineDataCopyPass();
 void registerTestAffineLoopParametricTilingPass();
 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
+void registerTestBufferizeStructuralOpsPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvVectorization();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
 void registerTestDialect(DialectRegistry &);
 void registerTestDynamicPipelinePass();

 #if MLIR_CUDA_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToCubinPass();
 #endif
 #if MLIR_ROCM_CONVERSIONS_ENABLED
   registerTestConvertGPUKernelToHsacoPass();
 #endif
   registerTestAffineLoopParametricTilingPass();
   registerTestBufferPlacementPreparationPass();
+  registerTestBufferizeStructuralOpsPass();
   registerTestDominancePass();
   registerTestDynamicPipelinePass();
   registerTestFunc();
   registerTestExpandTanhPass();
   registerTestGpuMemoryPromotionPass();
   registerTestInterfaces();
   registerTestLinalgCodegenStrategy();
   registerTestLinalgFusionTransforms();