3781b7905d
For cases where we can automatically construct the Attribute allow for more
user-friendly input. This is consistent with C++ builder generation as well
choice of which single builder to generate here (most
specialized/user-friendly).
Registration of attribute builders from more pythonic input is all Python side.
The downside is that
* extra checking to see if user provided a custom builder in op builders,
* the ODS attribute name is load bearing
upside is that
* easily change these/register dialect specific ones in downstream projects,
* adding support/changing to different convenience builders are all along with
the rest of the convenience functions in Python (and no additional changes
to tablegen file or recompilation needed);
Allow for both building with Attributes as well as raw inputs. This change
should therefore be backwards compatible as well as allow for avoiding
recreating Attribute where already available.
Differential Revision: https://reviews.llvm.org/D139568
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
# RUN: %PYTHON %s | FileCheck %s
|
|
|
|
from mlir.ir import *
|
|
import numpy as np
|
|
import mlir.dialects.func as func
|
|
import mlir.dialects.shape as shape
|
|
|
|
|
|
def run(f):
|
|
print("\nTEST:", f.__name__)
|
|
f()
|
|
return f
|
|
|
|
|
|
# CHECK-LABEL: TEST: testConstShape
|
|
@run
|
|
def testConstShape():
|
|
with Context() as ctx, Location.unknown():
|
|
module = Module.create()
|
|
f32 = F32Type.get()
|
|
with InsertionPoint(module.body):
|
|
@func.FuncOp.from_py_func(
|
|
RankedTensorType.get((12, ShapedType.get_dynamic_size()), f32))
|
|
def const_shape_tensor(arg):
|
|
shape.ConstWitnessOp(False)
|
|
shape.ConstSizeOp(30)
|
|
shape.ConstSizeOp(IntegerAttr.get(IndexType.get(), 40))
|
|
shape.ConstShapeOp([1, 2])
|
|
return shape.ConstShapeOp(
|
|
DenseElementsAttr.get(
|
|
np.array([3, 4], dtype=np.int64), type=IndexType.get()))
|
|
|
|
# CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
|
|
# CHECK-DAG: shape.const_witness false
|
|
# CHECK-DAG: shape.const_size 30
|
|
# CHECK-DAG: shape.const_size 40
|
|
# CHECK-DAG: shape.const_shape [1, 2] : tensor<2xindex>
|
|
# CHECK-DAG: shape.const_shape [3, 4] : tensor<2xindex>
|
|
print(module)
|