Skip to content
This repository was archived by the owner on Apr 23, 2021. It is now read-only.

Add vector memref support for mlir-cpu-runner func return #103

Closed
wants to merge 1 commit into from
Closed

Add vector memref support for mlir-cpu-runner func return #103

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions lib/ExecutionEngine/MemRefUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@

#include "llvm/Support/Error.h"
#include <numeric>
#include <stdlib.h>

using namespace mlir;

Expand All @@ -45,12 +46,14 @@ allocMemRefDescriptor(Type type, bool allocateData = true,
return make_string_error("memref with dynamic shapes not supported");

auto elementType = memRefType.getElementType();
if (!elementType.isF32())
VectorType vectorType = elementType.dyn_cast<VectorType>();
if (!elementType.isF32() &&
!(vectorType && vectorType.getElementType().isF32()))
return make_string_error(
"memref with element other than f32 not supported");
"memref with element other than f32 or vector of f32 not supported");

auto *descriptor =
reinterpret_cast<StaticFloatMemRef *>(malloc(sizeof(StaticFloatMemRef)));
static_cast<StaticFloatMemRef *>(malloc(sizeof(StaticFloatMemRef)));
if (!allocateData) {
descriptor->data = nullptr;
return descriptor;
Expand All @@ -59,7 +62,18 @@ allocMemRefDescriptor(Type type, bool allocateData = true,
auto shape = memRefType.getShape();
int64_t size = std::accumulate(shape.begin(), shape.end(), 1,
std::multiplies<int64_t>());
descriptor->data = reinterpret_cast<float *>(malloc(sizeof(float) * size));
// Align vector of f32 to the vector size boundary (to the closest greater
// power of two if the former isn't a power of two).
if (vectorType) {
int64_t numElements = vectorType.getNumElements();
size *= numElements;
size_t alignment = llvm::PowerOf2Ceil(numElements * sizeof(float));
posix_memalign(reinterpret_cast<void **>(&descriptor->data), alignment,
size * sizeof(float));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it work on windows?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does it work on windows?

I don't have a setup to test this. But I hope someone who has MLIR working on Windows could confirm. There are other alternatives and we could use malloc and shift to align but the code in context becomes relatively messier since we'd have to keep track of it for freeing the memref's.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You won't be able to test it because the LLVM jit doesn't support windows right now, but I can confirm that posix_memalign is not available on windows. The closest thing would be _aligned_malloc.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't have a setup to test this. But I hope someone who has MLIR working on Windows could confirm

I kicked the Kokoro builds after asking :)

The windows build failed actually.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume you have access to the log from the results below, but here it is:


7297
       "T:\src\build\projects\mlir\test\check-mlir.vcxproj" (default target) (1) ->
7298
       "T:\src\build\projects\mlir\tools\mlir-cpu-runner\mlir-cpu-runner.vcxproj" (default target) (44) ->
7299
       "T:\src\build\projects\mlir\lib\ExecutionEngine\MLIRExecutionEngine.vcxproj" (default target) (98) ->
7300
       (ClCompile target) ->
7301
         t:\src\github\llvm\llvm\projects\mlir\lib\executionengine\memrefutils.cpp(71): error C3861: 'posix_memalign': identifier not found [T:\src\build\projects\mlir\lib\ExecutionEngine\MLIRExecutionEngine.vcxproj]
7302
7303

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! Would there be a way to write conditionally compiled code so that _aligned_malloc is used for a Windows build? (I didn't see any such examples in MLIR.)

} else {
descriptor->data = static_cast<float *>(malloc(sizeof(float) * size));
}

for (int64_t i = 0; i < size; ++i) {
descriptor->data[i] = initialValue;
}
Expand Down
4 changes: 4 additions & 0 deletions lib/Support/JitRunner.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,10 @@ static void printOneMemRef(Type t, void *val) {
auto shape = memRefType.getShape();
int64_t size = std::accumulate(shape.begin(), shape.end(), 1,
std::multiplies<int64_t>());
if (auto vectorType = memRefType.getElementType().dyn_cast<VectorType>()) {
size *= vectorType.getNumElements();
}

for (int64_t i = 0; i < size; ++i) {
llvm::outs() << reinterpret_cast<StaticFloatMemRef *>(val)->data[i] << ' ';
}
Expand Down
26 changes: 26 additions & 0 deletions test/mlir-cpu-runner/simple.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
// RUN: mlir-cpu-runner -e foo -init-value 1000 %s | FileCheck -check-prefix=NOMAIN %s
// RUN: mlir-cpu-runner %s -O3 | FileCheck %s
// RUN: mlir-cpu-runner -e affine -init-value 2.0 %s | FileCheck -check-prefix=AFFINE %s
// RUN: mlir-cpu-runner -e bar -init-value 2.0 %s | FileCheck -check-prefix=BAR %s
// RUN: mlir-cpu-runner -e large_vec_memref -init-value 2.0 %s | FileCheck -check-prefix=LARGE-VEC %s

// RUN: cp %s %t
// RUN: mlir-cpu-runner %t -dump-object-file | FileCheck %t
Expand Down Expand Up @@ -49,3 +51,27 @@ func @affine(%a : memref<32xf32>) -> memref<32xf32> {
return %a : memref<32xf32>
}
// AFFINE: 4.2{{0+}}e+01

func @bar(%a : memref<16xvector<4xf32>>) -> memref<16xvector<4xf32>> {
%c0 = constant 0 : index
%c1 = constant 1 : index

%u = load %a[%c0] : memref<16xvector<4xf32>>
%v = load %a[%c1] : memref<16xvector<4xf32>>
%w = addf %u, %v : vector<4xf32>
store %w, %a[%c0] : memref<16xvector<4xf32>>

return %a : memref<16xvector<4xf32>>
}
// BAR: 4.{{0+}}e+00 4.{{0+}}e+00 4.{{0+}}e+00 4.{{0+}}e+00 2.{{0+}}e+00
// BAR-NEXT: 4.{{0+}}e+00 4.{{0+}}e+00 4.{{0+}}e+00 4.{{0+}}e+00 2.{{0+}}e+00

func @large_vec_memref(%arg2: memref<128x128xvector<8xf32>>) -> memref<128x128xvector<8xf32>> {
%c0 = constant 0 : index
%c127 = constant 127 : index
%v = constant dense<42.0> : vector<8xf32>
store %v, %arg2[%c0, %c0] : memref<128x128xvector<8xf32>>
store %v, %arg2[%c127, %c127] : memref<128x128xvector<8xf32>>
return %arg2 : memref<128x128xvector<8xf32>>
}
// LARGE-VEC: 4.200000e+01