[Mlir-commits] [mlir] 3f2891d - [mlir] Add python test for shape dialect

[Jacques Pienaar]

Author: Jacques Pienaar
Date: 2021-05-13T09:13:47-07:00
New Revision: 3f2891db6dd5684ee743055c0e86d0d3dd66c90b

URL: https://github.com/llvm/llvm-project/commit/3f2891db6dd5684ee743055c0e86d0d3dd66c90b
DIFF: https://github.com/llvm/llvm-project/commit/3f2891db6dd5684ee743055c0e86d0d3dd66c90b.diff

LOG: [mlir] Add python test for shape dialect

Add basic test for shape.const_shape op as start.

Differential Revision: https://reviews.llvm.org/D102341

Added: 
    mlir/test/python/dialects/shape.py

Modified: 
    

Removed: 
    


################################################################################
diff  --git a/mlir/test/python/dialects/shape.py b/mlir/test/python/dialects/shape.py
new file mode 100644
index 0000000000000..1772026cb9b92
--- /dev/null
+++ b/mlir/test/python/dialects/shape.py
@@ -0,0 +1,32 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir.ir import *
+import numpy as np
+import mlir.dialects.builtin as builtin
+import mlir.dialects.shape as shape
+
+
+def run(f):
+  print("\nTEST:", f.__name__)
+  f()
+  return f
+
+
+# CHECK-LABEL: TEST: testConstShape
+ at run
+def testConstShape():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f32 = F32Type.get()
+    indexT = IndexType.get()
+    with InsertionPoint(module.body):
+      @builtin.FuncOp.from_py_func(
+          RankedTensorType.get((12, -1), f32))
+      def const_shape_tensor(arg):
+        return shape.ConstShapeOp(RankedTensorType.get((2,), indexT),
+                                  DenseElementsAttr.get(np.array([10, 20])))
+
+    # CHECK-LABEL: func @const_shape_tensor(%arg0: tensor<12x?xf32>)
+    # CHECK: shape.const_shape [10, 20] : tensor<2xindex>
+    print(module)
+