[MLIR][Linalg] Vectorization Fail with reduction

MLS · June 28, 2024, 11:40am

Such linalg.generic can not be vetorized.

#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map14 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>

func.func @generic_1_12_197_197_1_12_197_1_f32_547(%arg0: memref<465708xf32>, %arg1: memref<2364xf32>) attributes {global_kernel = ""} {
    %expand_shape = memref.expand_shape %arg0 [[0, 1, 2, 3]] : memref<465708xf32> into memref<1x12x197x197xf32>
    %expand_shape_0 = memref.expand_shape %arg1 [[0, 1, 2, 3]] : memref<2364xf32> into memref<1x12x197x1xf32>
    linalg.generic {indexing_maps = [#map1, #map14], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%expand_shape : memref<1x12x197x197xf32>) outs(%expand_shape_0 : memref<1x12x197x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %0 = arith.addf %in, %out : f32
      linalg.yield %0 : f32
    }
    return
}

module attributes {transform.with_named_sequence} {
    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
      %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
      %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op
      transform.yield
    }
}

It’s report:

mlir-opt: /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:465: mlir::AffineMap reindexIndexingMap(mlir::AffineMap): Assertion `res.getNumDims() == res.getNumResults() && "expected reindexed map with same number of dims and results"' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: ../build-llvm/bin/mlir-opt --transform-interpreter test.func.mlir
 #0 0x000055dbe341ca76 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /home/mls/MLIR/MLIR/llvm/lib/Support/Unix/Signals.inc:723:0
 #1 0x000055dbe341ce86 PrintStackTraceSignalHandler(void*) /home/mls/MLIR/MLIR/llvm/lib/Support/Unix/Signals.inc:798:0
 #2 0x000055dbe341a444 llvm::sys::RunSignalHandlers() /home/mls/MLIR/MLIR/llvm/lib/Support/Signals.cpp:105:0
 #3 0x000055dbe341c369 SignalHandler(int) /home/mls/MLIR/MLIR/llvm/lib/Support/Unix/Signals.inc:413:0
 #4 0x00007f142a4c0980 __restore_rt (/lib/x86_64-linux-gnu/libpthread.so.0+0x12980)
 #5 0x00007f1429171e87 raise /build/glibc-CVJwZb/glibc-2.27/signal/../sysdeps/unix/sysv/linux/raise.c:51:0
 #6 0x00007f14291737f1 abort /build/glibc-CVJwZb/glibc-2.27/stdlib/abort.c:81:0
 #7 0x00007f14291633fa __assert_fail_base /build/glibc-CVJwZb/glibc-2.27/assert/assert.c:89:0
 #8 0x00007f1429163472 (/lib/x86_64-linux-gnu/libc.so.6+0x30472)
 #9 0x000055dbe4ab33a9 reindexIndexingMap(mlir::AffineMap) /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:464:0
#10 0x000055dbe4ab7a0f vectorizeAsLinalgGeneric(mlir::RewriterBase&, VectorizationState&, mlir::linalg::LinalgOp, llvm::SmallVectorImpl<mlir::Value>&) /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:1318:0
#11 0x000055dbe4ac2ac4 auto mlir::linalg::vectorize(mlir::RewriterBase&, mlir::Operation*, llvm::ArrayRef<long>, llvm::ArrayRef<bool>, bool, bool)::'lambda'(auto)::operator()<mlir::linalg::LinalgOp>(auto) const /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:1724:0
#12 0x000055dbe4ac2b5c llvm::TypeSwitch<mlir::Operation*, mlir::LogicalResult>& llvm::TypeSwitch<mlir::Operation*, mlir::LogicalResult>::Case<mlir::linalg::LinalgOp, mlir::linalg::vectorize(mlir::RewriterBase&, mlir::Operation*, llvm::ArrayRef<long>, llvm::ArrayRef<bool>, bool, bool)::'lambda'(auto)>(mlir::linalg::vectorize(mlir::RewriterBase&, mlir::Operation*, llvm::ArrayRef<long>, llvm::ArrayRef<bool>, bool, bool)::'lambda'(auto)&&) /home/mls/MLIR/MLIR/llvm/include/llvm/ADT/TypeSwitch.h:122:0
#13 0x000055dbe4aba668 mlir::linalg::vectorize(mlir::RewriterBase&, mlir::Operation*, llvm::ArrayRef<long>, llvm::ArrayRef<bool>, bool, bool) /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:1724:0
#14 0x000055dbe46e501a (anonymous namespace)::VectorizationPattern::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&) const /home/mls/MLIR/MLIR/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp:3025:0
#15 0x000055dbebe41c29 mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)::'lambda'()::operator()() const /home/mls/MLIR/MLIR/mlir/lib/Rewrite/PatternApplicator.cpp:206:0
#16 0x000055dbebe425dd void llvm::function_ref<void ()>::callback_fn<mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>)::'lambda'()>(long) /home/mls/MLIR/MLIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:0
#17 0x000055dbe343d5a8 llvm::function_ref<void ()>::operator()() const /home/mls/MLIR/MLIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:68:0
#18 0x000055dbebe45564 void mlir::MLIRContext::executeAction<mlir::ApplyPatternAction, mlir::Pattern const&>(llvm::function_ref<void ()>, llvm::ArrayRef<mlir::IRUnit>, mlir::Pattern const&) /home/mls/MLIR/MLIR/mlir/include/mlir/IR/MLIRContext.h:276:0
#19 0x000055dbebe4237c mlir::PatternApplicator::matchAndRewrite(mlir::Operation*, mlir::PatternRewriter&, llvm::function_ref<bool (mlir::Pattern const&)>, llvm::function_ref<void (mlir::Pattern const&)>, llvm::function_ref<mlir::LogicalResult (mlir::Pattern const&)>) /home/mls/MLIR/MLIR/mlir/lib/Rewrite/PatternApplicator.cpp:227:0
#20 0x000055dbe8229885 (anonymous namespace)::GreedyPatternRewriteDriver::processWorklist() /home/mls/MLIR/MLIR/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:474:0
#21 0x000055dbe822a880 (anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::'lambda2'()::operator()() const /home/mls/MLIR/MLIR/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp:693:0
#22 0x000055dbe822b99f void llvm::function_ref<void ()>::callback_fn<(anonymous namespace)::RegionPatternRewriteDriver::simplify(bool*) &&::'lambda2'()>(long) /home/mls/MLIR/MLIR/llvm/include/llvm/ADT/STLFunctionalExtras.h:46:0
...

github.com

llvm/llvm-project/blob/f1faba25433c971f024dd8a29da14020246e89ec/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp#L471-L478


      
          static AffineMap reindexIndexingMap(AffineMap map) {
            assert(map.isProjectedPermutation(/*allowZeroInResults=*/true) &&
                   "expected projected permutation");
            auto res = compressUnusedDims(map);
            assert(res.getNumDims() == res.getNumResults() &&
                   "expected reindexed map with same number of dims and results");
            return res;
          }

I know the error is because the compressed dimension is not equal to the result dimension, but I don’t know how to fix it.

dcaballe · June 28, 2024, 4:47pm

That seems to work for me with ToT:

mlir-opt test.mlir -transform-interpreter -mlir-print-ir-before-all
// -----// IR Dump Before InterpreterPass (transform-interpreter) //----- //
#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, 0)>
module {
  func.func @generic_1_12_197_197_1_12_197_1_f32_547(%arg0: memref<465708xf32>, %arg1: memref<2364xf32>) attributes {global_kernel = ""} {
    %expand_shape = memref.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 12, 197, 197] : memref<465708xf32> into memref<1x12x197x197xf32>
    %expand_shape_0 = memref.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 12, 197, 1] : memref<2364xf32> into memref<1x12x197x1xf32>
    linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%expand_shape : memref<1x12x197x197xf32>) outs(%expand_shape_0 : memref<1x12x197x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %0 = arith.addf %in, %out : f32
      linalg.yield %0 : f32
    }
    return
  }
  module attributes {transform.with_named_sequence} {
    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
      %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
      %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op
      transform.yield
    }
  }
}


#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
module {
  func.func @generic_1_12_197_197_1_12_197_1_f32_547(%arg0: memref<465708xf32>, %arg1: memref<2364xf32>) attributes {global_kernel = ""} {
    %cst = arith.constant 0.000000e+00 : f32
    %c0 = arith.constant 0 : index
    %expand_shape = memref.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 12, 197, 197] : memref<465708xf32> into memref<1x12x197x197xf32>
    %expand_shape_0 = memref.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 12, 197, 1] : memref<2364xf32> into memref<1x12x197x1xf32>
    %0 = vector.transfer_read %expand_shape[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true, true]} : memref<1x12x197x197xf32>, vector<1x12x197x197xf32>
    %1 = vector.transfer_read %expand_shape_0[%c0, %c0, %c0, %c0], %cst {in_bounds = [true, true, true], permutation_map = #map} : memref<1x12x197x1xf32>, vector<1x12x197xf32>
    %2 = vector.multi_reduction <add>, %0, %1 [3] : vector<1x12x197x197xf32> to vector<1x12x197xf32>
    %3 = vector.broadcast %2 : vector<1x12x197xf32> to vector<1x1x12x197xf32>
    %4 = vector.transpose %3, [1, 2, 3, 0] : vector<1x1x12x197xf32> to vector<1x12x197x1xf32>
    vector.transfer_write %4, %expand_shape_0[%c0, %c0, %c0, %c0] {in_bounds = [true, true, true, true]} : vector<1x12x197x1xf32>, memref<1x12x197x1xf32>
    return
  }
  module attributes {transform.with_named_sequence} {
    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
      %0 = transform.structured.match ops{["linalg.generic"]} in %arg0 : (!transform.any_op) -> !transform.any_op
      %1 = transform.get_parent_op %0 {isolated_from_above} : (!transform.any_op) -> !transform.any_op
      %2 = transform.structured.vectorize_children_and_apply_patterns %1 : (!transform.any_op) -> !transform.any_op
      transform.yield
    }
  }
}

MLS · June 29, 2024, 4:41am

I pulled the latest code from the community, but the problem persists.

ftynse · June 29, 2024, 6:51pm

Please provide full instructions for reproducing the issue. (Also, it seems better suited for a Github issue, as the message indicates, than a forum post.)

Topic		Replies	Views
How to efficiently vectorize linalg.generic reduction op? MLIR	2	529	January 9, 2024
Understanding Vectorization Failure with transform.structured.vectorize MLIR mlir	4	288	June 25, 2024
Vectorization failure MLIR	0	218	August 2, 2023
Vectorization of convolution op MLIR	5	945	February 23, 2022
Linalg.generic parallel vs. reduction MLIR	4	871	October 18, 2021

[MLIR][Linalg] Vectorization Fail with reduction

Related topics