dd473f1dd1
https://github.com/llvm/llvm-project/pull/68853 enabled a lot of nice cleanup. Note, I made sure each of the touched extensions had tests.
45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
import mlir.dialects.func as func
|
|
import mlir.dialects.arith as arith
|
|
import mlir.dialects.affine as affine
|
|
import mlir.dialects.memref as memref
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testAffineStoreOp
|
|
@run
|
|
def testAffineStoreOp():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
with InsertionPoint(module.body):
|
|
f32 = F32Type.get()
|
|
index_type = IndexType.get()
|
|
memref_type_out = MemRefType.get([12, 12], f32)
|
|
|
|
# CHECK: func.func @affine_store_test(%[[ARG0:.*]]: index) -> memref<12x12xf32> {
|
|
@func.FuncOp.from_py_func(index_type)
|
|
def affine_store_test(arg0):
|
|
# CHECK: %[[O_VAR:.*]] = memref.alloc() : memref<12x12xf32>
|
|
mem = memref.AllocOp(memref_type_out, [], []).result
|
|
|
|
d0 = AffineDimExpr.get(0)
|
|
s0 = AffineSymbolExpr.get(0)
|
|
map = AffineMap.get(1, 1, [s0 * 3, d0 + s0 + 1])
|
|
|
|
# CHECK: %[[A1:.*]] = arith.constant 2.100000e+00 : f32
|
|
a1 = arith.ConstantOp(f32, 2.1)
|
|
|
|
# CHECK: affine.store %[[A1]], %alloc[symbol(%[[ARG0]]) * 3, %[[ARG0]] + symbol(%[[ARG0]]) + 1] : memref<12x12xf32>
|
|
affine.AffineStoreOp(a1, mem, indices=[arg0, arg0], map=map)
|
|
|
|
return mem
|
|
|
|
print(module)
|