This is an archive of the discontinued LLVM Phabricator instance.

Differential D110444

[mlir][sparse] sampled matrix multiplication fusion test
ClosedPublic

Authored by aartbik on Sep 24 2021, 1:37 PM.

Details

Reviewers
bixia
ThomasRaoux
mravishankar
Commits
rG06e2a0684e52: [mlir][sparse] sampled matrix multiplication fusion test
Summary

This integration tests runs a fused and non-fused version of
sampled matrix multiplication. Both should eventually have the
same performance!

NOTE: relies on pending tensor.init fix!

Diff Detail

Event Timeline

aartbik created this revision.Sep 24 2021, 1:37 PM
Herald added subscribers: wenzhicui, wrengr, Chia-hungDuan and 19 others. · View Herald TranscriptSep 24 2021, 1:37 PM
aartbik requested review of this revision.Sep 24 2021, 1:37 PM
Herald added a project: Restricted Project. · View Herald TranscriptSep 24 2021, 1:37 PM
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
aartbik added reviewers: bixia, ThomasRaoux, mravishankar.Sep 24 2021, 1:39 PM
aartbik added a parent revision: D110442: [mlir][sparse] preserve zero-initialization for materializing buffers.
Harbormaster completed remote builds in B125637: Diff 374947.Sep 24 2021, 1:43 PM
bixia accepted this revision.Sep 27 2021, 9:01 AM
This revision is now accepted and ready to land.Sep 27 2021, 9:01 AM
aartbik updated this revision to Diff 375349.Sep 27 2021, 11:24 AM

rebased with parent in main

Harbormaster completed remote builds in B125933: Diff 375349.Sep 27 2021, 11:48 AM
Closed by commit rG06e2a0684e52: [mlir][sparse] sampled matrix multiplication fusion test (authored by aartbik). · Explain WhySep 27 2021, 11:50 AM
This revision was automatically updated to reflect the committed changes.
aartbik added a commit: rG06e2a0684e52: [mlir][sparse] sampled matrix multiplication fusion test.

Revision Contents

PathSize
mlir/
test/
Integration/
Dialect/
SparseTensor/
CPU/
168 lines

Diff 375350

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir

  • This file was added.
// RUN: mlir-opt %s \
// RUN: --linalg-generalize-named-ops --linalg-fuse-elementwise-ops \
// RUN: --sparsification --sparse-tensor-conversion \
// RUN: --linalg-bufferize --convert-linalg-to-loops \
// RUN: --convert-vector-to-scf --convert-scf-to-std \
// RUN: --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
// RUN: --std-bufferize --finalizing-bufferize --lower-affine \
// RUN: --convert-vector-to-llvm --convert-memref-to-llvm --convert-math-to-llvm \
// RUN: --convert-std-to-llvm --reconcile-unrealized-casts | \
// RUN: mlir-cpu-runner \
// RUN: -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s
//
// Do the same run, but now with SIMDization as well.
// This should not change the outcome.
//
// RUN: mlir-opt %s \
// RUN: --linalg-generalize-named-ops --linalg-fuse-elementwise-ops \
// RUN: --sparsification="vectorization-strategy=2 vl=8" --sparse-tensor-conversion \
// RUN: --linalg-bufferize --convert-linalg-to-loops \
// RUN: --convert-vector-to-scf --convert-scf-to-std \
// RUN: --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
// RUN: --std-bufferize --finalizing-bufferize --lower-affine \
// RUN: --convert-vector-to-llvm --convert-memref-to-llvm --convert-math-to-llvm \
// RUN: --convert-std-to-llvm --reconcile-unrealized-casts | \
// RUN: mlir-cpu-runner \
// RUN: -e entry -entry-point-result=void \
// RUN: -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
// RUN: FileCheck %s
#SM = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
#trait_sampled_dense_dense = {
indexing_maps = [
affine_map<(i,j,k) -> (i,j)>, // S
affine_map<(i,j,k) -> (i,k)>, // A
affine_map<(i,j,k) -> (k,j)>, // B
affine_map<(i,j,k) -> (i,j)> // X (out)
],
iterator_types = ["parallel", "parallel", "reduction"],
doc = "X(i,j) += S(i,j) SUM_k A(i,k) B(k,j)"
}
#trait_matmul = {
indexing_maps = [
affine_map<(d0, d1, d2) -> (d1, d0)>,
affine_map<(d0, d1, d2) -> (d0, d2)>,
affine_map<(d0, d1, d2) -> (d1, d2)>
],
iterator_types = ["reduction", "parallel", "parallel"]
}
#trait_scale = {
indexing_maps = [
affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>,
affine_map<(d0, d1) -> (d0, d1)>
],
iterator_types = ["parallel", "parallel"]
}
//
// Integration test for sampled dense dense matmul fusion.
//
module {
//
// A kernel that computes a direct sampled matrix matrix multiplication.
//
func @sampled_dd(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64> {
%d = constant 0.0 : f64
%0 = linalg.init_tensor [8, 8] : tensor<8x8xf64>
%1 = linalg.fill(%d, %0) : f64, tensor<8x8xf64> -> tensor<8x8xf64>
%2 = linalg.generic #trait_sampled_dense_dense
ins(%args, %arga, %argb: tensor<8x8xf64, #SM>,
tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1: tensor<8x8xf64>) {
^bb(%s: f64, %a: f64, %b: f64, %x: f64):
%p = mulf %a, %b : f64
%q = mulf %s, %p : f64
%r = addf %x, %q : f64
linalg.yield %r : f64
} -> tensor<8x8xf64>
return %2 : tensor<8x8xf64>
}
//
// A kernel that computes an unfused sampled matrix matrix multiplication.
//
func @sampled_dd_unfused(%args: tensor<8x8xf64, #SM>,
%arga: tensor<8x8xf64>,
%argb: tensor<8x8xf64>) -> tensor<8x8xf64> {
%d = constant 0.0 : f64
%0 = linalg.init_tensor [8, 8] : tensor<8x8xf64>
%1 = linalg.fill(%d, %0) : f64, tensor<8x8xf64> -> tensor<8x8xf64>
%2 = linalg.generic #trait_matmul
ins(%arga, %argb : tensor<8x8xf64>, tensor<8x8xf64>)
outs(%1 : tensor<8x8xf64>) {
^bb0(%a: f64, %b: f64, %x: f64):
%p = mulf %a, %b : f64
%q = addf %x, %p : f64
linalg.yield %q : f64
} -> tensor<8x8xf64>
%3 = linalg.init_tensor [8, 8] : tensor<8x8xf64>
%4 = linalg.fill(%d, %3) : f64, tensor<8x8xf64> -> tensor<8x8xf64>
%5 = linalg.generic #trait_scale
ins(%2, %args : tensor<8x8xf64>, tensor<8x8xf64, #SM>)
outs(%4 : tensor<8x8xf64>) {
^bb0(%t: f64, %s: f64, %x: f64):
%r = mulf %t, %s : f64
linalg.yield %r : f64
} -> tensor<8x8xf64>
return %5 : tensor<8x8xf64>
}
//
// Main driver.
//
func @entry() {
%d0 = constant 0.0 : f64
%c0 = constant 0 : index
%t = constant sparse<[[0, 0], [7,7]], [1.0, 2.0]>
: tensor<8x8xf64>
%s = sparse_tensor.convert %t
: tensor<8x8xf64> to tensor<8x8xf64, #SM>
%a = constant dense<3.0> : tensor<8x8xf64>
%b = constant dense<4.0> : tensor<8x8xf64>
// Call the kernels.
%0 = call @sampled_dd(%s, %a, %b)
: (tensor<8x8xf64, #SM>,
tensor<8x8xf64>, tensor<8x8xf64>) -> tensor<8x8xf64>
%1 = call @sampled_dd_unfused(%s, %a, %b)
: (tensor<8x8xf64, #SM>,
tensor<8x8xf64>, tensor<8x8xf64>) -> tensor<8x8xf64>
// Verify the outputs.
//
// CHECK: ( ( 96, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 192 ) )
//
// CHECK: ( ( 96, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 0 ),
// CHECK-SAME: ( 0, 0, 0, 0, 0, 0, 0, 0 ), ( 0, 0, 0, 0, 0, 0, 0, 192 ) )
//
%m0 = memref.buffer_cast %0 : memref<8x8xf64>
%m1 = memref.buffer_cast %1 : memref<8x8xf64>
%v0 = vector.transfer_read %m0[%c0, %c0], %d0
: memref<8x8xf64>, vector<8x8xf64>
%v1 = vector.transfer_read %m1[%c0, %c0], %d0
: memref<8x8xf64>, vector<8x8xf64>
vector.print %v0 : vector<8x8xf64>
vector.print %v1 : vector<8x8xf64>
return
}
}