175 lines · c
1// Split the MLIR string: this will produce %t/input.mlir2// RUN: split-file %s %t3 4// Compile the MLIR file to LLVM:5// RUN: mlir-opt %t/input.mlir \6// RUN: -lower-affine -convert-scf-to-cf -finalize-memref-to-llvm \7// RUN: -convert-func-to-llvm -convert-arith-to-llvm -convert-cf-to-llvm \8// RUN: -reconcile-unrealized-casts \9// RUN: | mlir-translate --mlir-to-llvmir -o %t.ll10 11// Generate an object file for the MLIR code12// RUN: llc %t.ll -o %t.o -filetype=obj13 14// Compile the current C file and link it to the MLIR code:15// RUN: "%host_cc" %s %t.o -o %t.exe16 17// Exec18// RUN: %t.exe | FileCheck %s19 20/* MLIR_BEGIN21//--- input.mlir22// Performs: arg0[i, j] = arg0[i, j] + arg1[i, j]23func.func private @add_memref(%arg0: memref<?x?xf64>, %arg1: memref<?x?xf64>) -> i6424 attributes {llvm.emit_c_interface} {25 %c0 = arith.constant 0 : index26 %c1 = arith.constant 1 : index27 %dimI = memref.dim %arg0, %c0 : memref<?x?xf64>28 %dimJ = memref.dim %arg0, %c1 : memref<?x?xf64>29 affine.for %i = 0 to %dimI {30 affine.for %j = 0 to %dimJ {31 %load0 = memref.load %arg0[%i, %j] : memref<?x?xf64>32 %load1 = memref.load %arg1[%i, %j] : memref<?x?xf64>33 %add = arith.addf %load0, %load1 : f6434 affine.store %add, %arg0[%i, %j] : memref<?x?xf64>35 }36 }37 %c42 = arith.constant 42 : i6438 return %c42 : i6439}40 41//--- end_input.mlir42 43MLIR_END */44 45#include <stdint.h>46#include <stdio.h>47 48// Define the API for the MLIR function, see49// https://mlir.llvm.org/docs/TargetLLVMIR/#calling-conventions for details.50//51// The function takes two 2D memref, the signature in MLIR LLVM dialect will be:52// llvm.func @add_memref(53// // First Memref (%arg0)54// %allocated_ptr0: !llvm.ptr<f64>, %aligned_ptr0: !llvm.ptr<f64>,55// %offset0: i64, %size0_d0: i64, %size0_d1: i64, %stride0_d0: i64,56// %stride0_d1: i64,57// // Second Memref (%arg1)58// %allocated_ptr1: !llvm.ptr<f64>, %aligned_ptr1: !llvm.ptr<f64>,59// %offset1: i64, %size1_d0: i64, %size1_d1: i64, %stride1_d0: i64,60// %stride1_d1: i64,61//62long long add_memref(double *allocated_ptr0, double *aligned_ptr0,63 intptr_t offset0, intptr_t size0_d0, intptr_t size0_d1,64 intptr_t stride0_d0, intptr_t stride0_d1,65 // Second Memref (%arg1)66 double *allocated_ptr1, double *aligned_ptr1,67 intptr_t offset1, intptr_t size1_d0, intptr_t size1_d1,68 intptr_t stride1_d0, intptr_t stride1_d1);69 70// The llvm.emit_c_interface will also trigger emission of another wrapper:71// llvm.func @_mlir_ciface_add_memref(72// %arg0: !llvm.ptr<struct<(ptr<f64>, ptr<f64>, i64,73// array<2 x i64>, array<2 x i64>)>>,74// %arg1: !llvm.ptr<struct<(ptr<f64>, ptr<f64>, i64,75// array<2 x i64>, array<2 x i64>)>>)76// -> i6477typedef struct {78 double *allocated;79 double *aligned;80 intptr_t offset;81 intptr_t size[2];82 intptr_t stride[2];83} memref_2d_descriptor;84long long _mlir_ciface_add_memref(memref_2d_descriptor *arg0,85 memref_2d_descriptor *arg1);86 87#define N 488#define M 889double arg0[N][M];90double arg1[N][M];91 92void dump() {93 for (int i = 0; i < N; i++) {94 printf("[");95 for (int j = 0; j < M; j++)96 printf("%d,\t", (int)arg0[i][j]);97 printf("] [");98 for (int j = 0; j < M; j++)99 printf("%d,\t", (int)arg1[i][j]);100 printf("]\n");101 }102}103 104int main() {105 int count = 0;106 for (int i = 0; i < N; i++) {107 for (int j = 0; j < M; j++) {108 arg0[i][j] = count++;109 arg1[i][j] = count++;110 }111 }112 printf("Before:\n");113 dump();114 // clang-format off115 // CHECK-LABEL: Before:116 // CHECK: [0, 2, 4, 6, 8, 10, 12, 14, ] [1, 3, 5, 7, 9, 11, 13, 15, ]117 // CHECK: [16, 18, 20, 22, 24, 26, 28, 30, ] [17, 19, 21, 23, 25, 27, 29, 31, ]118 // CHECK: [32, 34, 36, 38, 40, 42, 44, 46, ] [33, 35, 37, 39, 41, 43, 45, 47, ]119 // CHECK: [48, 50, 52, 54, 56, 58, 60, 62, ] [49, 51, 53, 55, 57, 59, 61, 63, ]120 // clang-format on121 122 // Call into MLIR.123 long long result = add_memref((double *)arg0, (double *)arg0, 0, N, M, M, 0,124 //125 (double *)arg1, (double *)arg1, 0, N, M, M, 0);126 127 // CHECK-LABEL: Result:128 // CHECK: 42129 printf("Result: %d\n", (int)result);130 131 printf("After:\n");132 dump();133 134 // clang-format off135 // CHECK-LABEL: After:136 // CHECK: [1, 5, 9, 13, 17, 21, 25, 29, ] [1, 3, 5, 7, 9, 11, 13, 15, ] 137 // CHECK: [33, 37, 41, 45, 49, 53, 57, 61, ] [17, 19, 21, 23, 25, 27, 29, 31, ]138 // CHECK: [65, 69, 73, 77, 81, 85, 89, 93, ] [33, 35, 37, 39, 41, 43, 45, 47, ]139 // CHECK: [97, 101, 105, 109, 113, 117, 121, 125, ] [49, 51, 53, 55, 57, 59, 61, 63, ]140 // clang-format on141 142 // Reset the input and re-apply the same function use the C API wrapper.143 count = 0;144 for (int i = 0; i < N; i++) {145 for (int j = 0; j < M; j++) {146 arg0[i][j] = count++;147 arg1[i][j] = count++;148 }149 }150 151 // Call into MLIR.152 memref_2d_descriptor arg0_descriptor = {153 (double *)arg0, (double *)arg0, 0, N, M, M, 0};154 memref_2d_descriptor arg1_descriptor = {155 (double *)arg1, (double *)arg1, 0, N, M, M, 0};156 result = _mlir_ciface_add_memref(&arg0_descriptor, &arg1_descriptor);157 158 // CHECK-LABEL: Result2:159 // CHECK: 42160 printf("Result2: %d\n", (int)result);161 162 printf("After2:\n");163 dump();164 165 // clang-format off166 // CHECK-LABEL: After2:167 // CHECK: [1, 5, 9, 13, 17, 21, 25, 29, ] [1, 3, 5, 7, 9, 11, 13, 15, ] 168 // CHECK: [33, 37, 41, 45, 49, 53, 57, 61, ] [17, 19, 21, 23, 25, 27, 29, 31, ]169 // CHECK: [65, 69, 73, 77, 81, 85, 89, 93, ] [33, 35, 37, 39, 41, 43, 45, 47, ]170 // CHECK: [97, 101, 105, 109, 113, 117, 121, 125, ] [49, 51, 53, 55, 57, 59, 61, 63, ]171 // clang-format on172 173 return 0;174}175