[mlir][EDSC] Add divis and diviu and vector.extractelement

Reviewed By: ftynse

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

diff --git a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h
index eccb980..f654020 100644 (file)
--- a/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h
+++ b/mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h
@@ -25,6 +25,8 @@ using std_constant_float = ValueBuilder<ConstantFloatOp>;
 using std_constant_index = ValueBuilder<ConstantIndexOp>;
 using std_constant_int = ValueBuilder<ConstantIntOp>;
 using std_dealloc = OperationBuilder<DeallocOp>;
+using std_divis = ValueBuilder<SignedDivIOp>;
+using std_diviu = ValueBuilder<UnsignedDivIOp>;
 using std_dim = ValueBuilder<DimOp>;
 using std_extract_element = ValueBuilder<ExtractElementOp>;
 using std_im = ValueBuilder<ImOp>;

diff --git a/mlir/include/mlir/Dialect/Vector/EDSC/Intrinsics.h b/mlir/include/mlir/Dialect/Vector/EDSC/Intrinsics.h
index 342f70a..f353262 100644 (file)
--- a/mlir/include/mlir/Dialect/Vector/EDSC/Intrinsics.h
+++ b/mlir/include/mlir/Dialect/Vector/EDSC/Intrinsics.h
@@ -19,6 +19,7 @@ using vector_contract = ValueBuilder<vector::ContractionOp>;
 using vector_insert = ValueBuilder<vector::InsertOp>;
 using vector_fma = ValueBuilder<vector::FMAOp>;
 using vector_extract = ValueBuilder<vector::ExtractOp>;
+using vector_extractelement = ValueBuilder<vector::ExtractElementOp>;
 using vector_matmul = ValueBuilder<vector::MatmulOp>;
 using vector_outerproduct = ValueBuilder<vector::OuterProductOp>;
 using vector_print = OperationBuilder<vector::PrintOp>;
@@ -26,9 +27,6 @@ using vector_transfer_read = ValueBuilder<vector::TransferReadOp>;
 using vector_transfer_write = OperationBuilder<vector::TransferWriteOp>;
 using vector_transpose = ValueBuilder<vector::TransposeOp>;
 using vector_type_cast = ValueBuilder<vector::TypeCastOp>;
-using vector_insert = ValueBuilder<vector::InsertOp>;
-using vector_fma = ValueBuilder<vector::FMAOp>;
-using vector_extract = ValueBuilder<vector::ExtractOp>;
 using vector_extract_slices = ValueBuilder<vector::ExtractSlicesOp>;
 using vector_insert_slices = ValueBuilder<vector::InsertSlicesOp>;
 using vector_extract_strided_slice =

diff --git a/mlir/test/EDSC/builder-api-test.cpp b/mlir/test/EDSC/builder-api-test.cpp
index 3eecf7e..3174881 100644 (file)
--- a/mlir/test/EDSC/builder-api-test.cpp
+++ b/mlir/test/EDSC/builder-api-test.cpp
@@ -419,6 +419,44 @@ TEST_FUNC(operator_and) {
   f.erase();
 }
 
+TEST_FUNC(divis_op_i32) {
+  using namespace edsc::op;
+  auto f = makeFunction("divis_op", {}, {});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  auto i32Type = builder.getI32Type();
+  std_divis(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
+
+  // clang-format off
+  // CHECK-LABEL: @divis_op
+  //   CHECK-DAG:     {{.*}} = constant 10
+  //   CHECK-DAG:     {{.*}} = constant 2
+  //  CHECK-NEXT:     {{.*}} = divi_signed
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
+TEST_FUNC(diviu_op_i32) {
+  using namespace edsc::op;
+  auto f = makeFunction("diviu_op", {}, {});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  auto i32Type = builder.getI32Type();
+  std_diviu(std_constant_int(10, i32Type), std_constant_int(2, i32Type));
+
+  // clang-format off
+  // CHECK-LABEL: @diviu_op
+  //   CHECK-DAG:     {{.*}} = constant 10
+  //   CHECK-DAG:     {{.*}} = constant 2
+  //  CHECK-NEXT:     {{.*}} = divi_unsigned
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
 TEST_FUNC(select_op_i32) {
   using namespace edsc::op;
   auto f32Type = FloatType::getF32(&globalContext());
@@ -924,6 +962,28 @@ TEST_FUNC(linalg_tensors_test) {
   f.erase();
 }
 
+TEST_FUNC(vector_extractelement_op_i32) {
+  using namespace edsc::op;
+  auto f = makeFunction("vector_extractelement_op", {}, {});
+
+  OpBuilder builder(f.getBody());
+  ScopedContext scope(builder, f.getLoc());
+  auto i32Type = builder.getI32Type();
+  auto vectorType = VectorType::get(/*shape=*/{8}, i32Type);
+  vector_extractelement(
+      i32Type, std_constant(vectorType, builder.getI32VectorAttr({10})),
+      std_constant_int(0, i32Type));
+
+  // clang-format off
+  // CHECK-LABEL: @vector_extractelement_op
+  //   CHECK-DAG:     {{.*}} = constant dense<10>
+  //   CHECK-DAG:     {{.*}} = constant 0
+  //  CHECK-NEXT:     {{.*}} = vector.extractelement
+  // clang-format on
+  f.print(llvm::outs());
+  f.erase();
+}
+
 // CHECK-LABEL: func @memref_vector_matmul_test(
 //  CHECK-SAME:   %[[A:.*]]: memref<?x?xvector<4x16xf32>>,
 //  CHECK-SAME:   %[[B:.*]]: memref<?x?xvector<16x8xf32>>,