This is an archive of the discontinued LLVM Phabricator instance.

Differential D138343

[mlir][sparse][vector] ensure loop peeling to remove vector masks works
ClosedPublic

Authored by aartbik on Nov 18 2022, 4:57 PM.

Details

Reviewers
nicolasvasilache
dcaballe
Peiming
bixia
qcolombet
Commits
rG914eff5a8ca9: [mlir][sparse][vector] ensure loop peeling to remove vector masks works

Diff Detail

Event Timeline

aartbik created this revision.Nov 18 2022, 4:57 PM
Herald added a project: Restricted Project. · View Herald TranscriptNov 18 2022, 4:57 PM
Herald added subscribers: jsetoain, Moerafaat, anlunx and 21 others. · View Herald Transcript
aartbik requested review of this revision.Nov 18 2022, 4:57 PM
Herald added a reviewer: nicolasvasilache. · View Herald TranscriptNov 18 2022, 4:57 PM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
aartbik added a parent revision: D138236: [mlir][sparse] introduce vectorization pass for sparse loops.Nov 18 2022, 4:57 PM
aartbik added reviewers: dcaballe, Peiming, bixia, qcolombet.
Peiming accepted this revision.Nov 18 2022, 5:03 PM
Peiming added inline comments.
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
2190

The b

This revision is now accepted and ready to land.Nov 18 2022, 5:03 PM
aartbik added inline comments.Nov 18 2022, 5:09 PM
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
2190

Yeah, bazel change belongs to parent revision.
Will be rebased away before merging. Sorry about that.

Harbormaster completed remote builds in B198587: Diff 476637.Nov 18 2022, 5:36 PM
Peiming added inline comments.Nov 18 2022, 7:39 PM
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
2190

Okay, seems that you inferred the whole sentence from "The b" already ;-)

Sorry that I accidentally submit the draft.

aartbik marked an inline comment as done.Nov 21 2022, 4:20 PM
aartbik added inline comments.
utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
2190

We work so closely together that I literally just need one letter reviews ;-)

aartbik updated this revision to Diff 477026.Nov 21 2022, 4:24 PM
aartbik marked an inline comment as done.

rebased main

This revision was landed with ongoing or failed builds.Nov 21 2022, 4:51 PM
Closed by commit rG914eff5a8ca9: [mlir][sparse][vector] ensure loop peeling to remove vector masks works (authored by aartbik). · Explain Why
This revision was automatically updated to reflect the committed changes.
aartbik added a commit: rG914eff5a8ca9: [mlir][sparse][vector] ensure loop peeling to remove vector masks works.
Harbormaster completed remote builds in B198879: Diff 477026.Nov 22 2022, 12:06 AM

Revision Contents

PathSize
mlir/
test/
Dialect/
SparseTensor/
63 lines
utils/
bazel/
llvm-project-overlay/
mlir/
2 lines

Diff 476637

mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir

  • This file was added.
// RUN: mlir-opt %s --sparsification -cse -sparse-vectorization="vl=16" -scf-for-loop-peeling -canonicalize -cse | \
// RUN: FileCheck %s
#SparseVector = #sparse_tensor.encoding<{
dimLevelType = [ "compressed" ],
pointerBitWidth = 32,
indexBitWidth = 32
}>
#trait_mul_s = {
indexing_maps = [
affine_map<(i) -> (i)>, // a
affine_map<(i) -> (i)>, // b
affine_map<(i) -> (i)> // x (out)
],
iterator_types = ["parallel"],
doc = "x(i) = a(i) * b(i)"
}
// CHECK-DAG: #[[$map0:.*]] = affine_map<()[s0, s1] -> (s0 + ((-s0 + s1) floordiv 16) * 16)>
// CHECK-DAG: #[[$map1:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>
// CHECK-LABEL: func @mul_s
// CHECK-DAG: %[[c0:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index
// CHECK-DAG: %[[c16:.*]] = arith.constant 16 : index
// CHECK: %[[p:.*]] = memref.load %{{.*}}[%[[c0]]] : memref<?xi32>
// CHECK: %[[a:.*]] = arith.extui %[[p]] : i32 to i64
// CHECK: %[[q:.*]] = arith.index_cast %[[a]] : i64 to index
// CHECK: %[[r:.*]] = memref.load %{{.*}}[%[[c1]]] : memref<?xi32>
// CHECK: %[[b:.*]] = arith.extui %[[r]] : i32 to i64
// CHECK: %[[s:.*]] = arith.index_cast %[[b]] : i64 to index
// CHECK: %[[boundary:.*]] = affine.apply #[[$map0]]()[%[[q]], %[[s]]]
// CHECK: scf.for %[[i:.*]] = %[[q]] to %[[boundary]] step %[[c16]] {
// CHECK: %[[mask:.*]] = vector.constant_mask [16] : vector<16xi1>
// CHECK: %[[li:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xi32>, vector<16xi32>
// CHECK: %[[zi:.*]] = arith.extui %[[li]] : vector<16xi32> to vector<16xi64>
// CHECK: %[[la:.*]] = vector.load %{{.*}}[%[[i]]] : memref<?xf32>, vector<16xf32>
// CHECK: %[[lb:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
// CHECK: %[[m:.*]] = arith.mulf %[[la]], %[[lb]] : vector<16xf32>
// CHECK: vector.scatter %{{.*}}[%[[c0]]] [%[[zi]]], %[[mask]], %[[m]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
// CHECK: }
// CHECK: scf.for %[[i2:.*]] = %[[boundary]] to %[[s]] step %[[c16]] {
// CHECK: %[[sub:.*]] = affine.apply #[[$map1]](%[[i2]])[%[[s]]]
// CHECK: %[[mask2:.*]] = vector.create_mask %[[sub]] : vector<16xi1>
// CHECK: %[[li2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xi32>, vector<16xi1>, vector<16xi32> into vector<16xi32>
// CHECK: %[[zi2:.*]] = arith.extui %[[li2]] : vector<16xi32> to vector<16xi64>
// CHECK: %[[la2:.*]] = vector.maskedload %{{.*}}[%[[i2]]], %[[mask2]], %{{.*}} : memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
// CHECK: %[[lb2:.*]] = vector.gather %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %{{.*}} : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32> into vector<16xf32>
// CHECK: %[[m2:.*]] = arith.mulf %[[la2]], %[[lb2]] : vector<16xf32>
// CHECK: vector.scatter %{{.*}}[%[[c0]]] [%[[zi2]]], %[[mask2]], %[[m2]] : memref<1024xf32>, vector<16xi64>, vector<16xi1>, vector<16xf32>
// CHECK: }
// CHECK: return
//
func.func @mul_s(%arga: tensor<1024xf32, #SparseVector>, %argb: tensor<1024xf32>, %argx: tensor<1024xf32>) -> tensor<1024xf32> {
%0 = linalg.generic #trait_mul_s
ins(%arga, %argb: tensor<1024xf32, #SparseVector>, tensor<1024xf32>)
outs(%argx: tensor<1024xf32>) {
^bb(%a: f32, %b: f32, %x: f32):
%0 = arith.mulf %a, %b : f32
linalg.yield %0 : f32
} -> tensor<1024xf32>
return %0 : tensor<1024xf32>
}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

Show First 20 Lines • Show All 2,169 Lines • ▼ Show 20 Lines deps = [
":DialectUtils", ":DialectUtils",
":FuncDialect", ":FuncDialect",
":FuncTransforms", ":FuncTransforms",
":IR", ":IR",
":LLVMDialect", ":LLVMDialect",
":LinalgDialect", ":LinalgDialect",
":LinalgTransforms", ":LinalgTransforms",
":LinalgUtils", ":LinalgUtils",
":MathDialect",
":MemRefDialect", ":MemRefDialect",
":Pass", ":Pass",
":SCFDialect", ":SCFDialect",
":SCFTransforms", ":SCFTransforms",
":SparseTensorDialect", ":SparseTensorDialect",
":SparseTensorEnums", ":SparseTensorEnums",
":SparseTensorPassIncGen", ":SparseTensorPassIncGen",
":SparseTensorUtils", ":SparseTensorUtils",
":Support", ":Support",
":TensorDialect", ":TensorDialect",
":Transforms", ":Transforms",
":VectorDialect",
PeimingUnsubmitted
Not Done
ReplyInline Actions

The b

Peiming: The b
aartbikAuthorUnsubmitted
Done
ReplyInline Actions

Yeah, bazel change belongs to parent revision.
Will be rebased away before merging. Sorry about that.

aartbik: Yeah, bazel change belongs to parent revision. Will be rebased away before merging. Sorry about…
PeimingUnsubmitted
Done
ReplyInline Actions

Okay, seems that you inferred the whole sentence from "The b" already ;-)

Sorry that I accidentally submit the draft.

Peiming: Okay, seems that you inferred the whole sentence from "The b" already ;-) Sorry that I…
aartbikAuthorUnsubmitted
Done
ReplyInline Actions

We work so closely together that I literally just need one letter reviews ;-)

aartbik: We work so closely together that I literally just need one letter reviews ;-)
"//llvm:Support", "//llvm:Support",
], ],
) )
cc_library( cc_library(
name = "SparseTensorPipelines", name = "SparseTensorPipelines",
srcs = glob(["lib/Dialect/SparseTensor/Pipelines/*.cpp"]), srcs = glob(["lib/Dialect/SparseTensor/Pipelines/*.cpp"]),
hdrs = ["include/mlir/Dialect/SparseTensor/Pipelines/Passes.h"], hdrs = ["include/mlir/Dialect/SparseTensor/Pipelines/Passes.h"],
▲ Show 20 Lines • Show All 8,122 Lines • Show Last 20 Lines