[Mlir-commits] [mlir] 2d0ca6b - [mlir][test] Add expand-strided-metadata to lower-to-llvm

[Quentin Colombet]

Author: Quentin Colombet
Date: 2022-12-15T09:07:21Z
New Revision: 2d0ca6b68ce246fbc2280e6f9cd5f5f649931b56

URL: https://github.com/llvm/llvm-project/commit/2d0ca6b68ce246fbc2280e6f9cd5f5f649931b56
DIFF: https://github.com/llvm/llvm-project/commit/2d0ca6b68ce246fbc2280e6f9cd5f5f649931b56.diff

LOG: [mlir][test] Add expand-strided-metadata to lower-to-llvm

This patch adds the `ExpandStridedMetadataPass` to the LowerToLLVM named
pipeline.
This is required now that "complex" memref operations, like subviews,
need to be expanded before being lowered.

Differential Revision: https://reviews.llvm.org/D139841

Added: 
    mlir/test/Dialect/LLVM/lower-to-llvm-e2e.mlir

Modified: 
    mlir/test/lib/Dialect/LLVM/CMakeLists.txt
    mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp
    utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

Removed: 
    


################################################################################
diff  --git a/mlir/test/Dialect/LLVM/lower-to-llvm-e2e.mlir b/mlir/test/Dialect/LLVM/lower-to-llvm-e2e.mlir
new file mode 100644
index 0000000000000..998650a7aeceb
--- /dev/null
+++ b/mlir/test/Dialect/LLVM/lower-to-llvm-e2e.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-opt %s -o - -test-lower-to-llvm -cse -split-input-file | FileCheck %s
+// Note: We run CSE here to make the pattern matching more direct.
+
+// Check that we properly lower to llvm memref operations that require to be
+// expanded first, like `memref.subview`.
+func.func @subview(%0 : memref<64x4xf32, strided<[4, 1], offset: 0>>, %arg0 : index, %arg1 : index, %arg2 : index)
+-> memref<?x?xf32, strided<[?, ?], offset: ?>> {
+  // CHECK-LABEL: @subview
+  // CHECK-SAME: %[[BASE:[^:]*]]: !llvm.ptr<f32>,
+  // CHECK-SAME: %[[BASE_ALIGNED:[^:]*]]: !llvm.ptr<f32>,
+  // CHECK-SAME: %[[BASE_OFFSET:[^:]*]]: i64,
+  // CHECK-SAME: %[[BASE_STRIDE0:[^:]*]]: i64,
+  // CHECK-SAME: %[[BASE_STRIDE1:[^:]*]]: i64,
+  // CHECK-SAME: %[[BASE_SIZE0:[^:]*]]: i64,
+  // CHECK-SAME: %[[BASE_SIZE1:[^:]*]]: i64,
+  // CHECK-SAME: %[[ARG0:[^:]*]]: i64,
+  // CHECK-SAME: %[[ARG1:[^:]*]]: i64,
+  // CHECK-SAME: %[[ARG2:[^:]*]]: i64)
+  // CHECK-SAME: -> !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>
+
+  // CHECK-DAG: %[[STRIDE0:.*]] = llvm.mlir.constant(4 : index) : i64
+  // CHECK-DAG: %[[DESCSTRIDE0:.*]] = llvm.mul %[[ARG0]], %[[STRIDE0]] : i64
+  // CHECK-DAG: %[[OFF2:.*]] = llvm.add %[[DESCSTRIDE0]], %[[ARG1]] : i64
+  // CHECK-DAG: %[[DESC:.*]] = llvm.mlir.undef : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+
+  // Base address and algined address.
+  // CHECK-DAG: %[[DESC0:.*]] = llvm.insertvalue %[[BASE]], %[[DESC]][0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+  // CHECK-DAG: %[[DESC1:.*]] = llvm.insertvalue %[[BASE_ALIGNED]], %[[DESC0]][1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+
+  // Offset.
+  // CHECK: %[[DESC2:.*]] = llvm.insertvalue %[[OFF2]], %[[DESC1]][2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+  // Size 0.
+  // CHECK: %[[DESC3:.*]] = llvm.insertvalue %[[ARG0]], %[[DESC2]][3, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+  // Stride 0 == 4 * %arg0.
+  // CHECK: %[[DESC4:.*]] = llvm.insertvalue %[[DESCSTRIDE0]], %[[DESC3]][4, 0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+  // Size 1.
+  // CHECK: %[[DESC5:.*]] = llvm.insertvalue %[[ARG1]], %[[DESC4]][3, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+  // Stride 1 == 1 * %arg1.
+  // CHECK: %[[DESC6:.*]] = llvm.insertvalue %[[ARG1]], %[[DESC5]][4, 1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64, array<2 x i64>, array<2 x i64>)>
+
+  %1 = memref.subview %0[%arg0, %arg1][%arg0, %arg1][%arg0, %arg1] :
+    memref<64x4xf32, strided<[4, 1], offset: 0>>
+  to memref<?x?xf32, strided<[?, ?], offset: ?>>
+  return %1 : memref<?x?xf32, strided<[?, ?], offset: ?>>
+}

diff  --git a/mlir/test/lib/Dialect/LLVM/CMakeLists.txt b/mlir/test/lib/Dialect/LLVM/CMakeLists.txt
index 5c9fe8410d97c..1ddb93adafbf1 100644
--- a/mlir/test/lib/Dialect/LLVM/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/LLVM/CMakeLists.txt
@@ -15,6 +15,7 @@ add_mlir_library(MLIRLLVMTestPasses
   MLIRLinalgTransforms
   MLIRMathToLLVM
   MLIRMemRefToLLVM
+  MLIRMemRefTransforms
   MLIRPass
   MLIRReconcileUnrealizedCasts
   MLIRSCFToControlFlow

diff  --git a/mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp b/mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp
index 503817f5a42d7..f863240f34657 100644
--- a/mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp
+++ b/mlir/test/lib/Dialect/LLVM/TestLowerToLLVM.cpp
@@ -24,6 +24,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
@@ -68,6 +69,10 @@ void buildTestLowerToLLVM(OpPassManager &pm,
           options.reassociateFPReductions)));
   // Convert Math to LLVM (always needed).
   pm.addNestedPass<func::FuncOp>(createConvertMathToLLVMPass());
+  // Expand complicated MemRef operations before lowering them.
+  pm.addPass(memref::createExpandStridedMetadataPass());
+  // The expansion may create affine expressions. Get rid of them.
+  pm.addPass(createLowerAffinePass());
   // Convert MemRef to LLVM (always needed).
   pm.addPass(createMemRefToLLVMConversionPass());
   // Convert Func to LLVM (always needed).

diff  --git a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
index a400c63efbee9..aa262f2a1aa79 100644
--- a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
@@ -599,6 +599,7 @@ cc_library(
         "//mlir:LinalgTransforms",
         "//mlir:MathToLLVM",
         "//mlir:MemRefToLLVM",
+        "//mlir:MemRefTransforms",
         "//mlir:Pass",
         "//mlir:ReconcileUnrealizedCasts",
         "//mlir:SCFToControlFlow",